Update midi_model.py

midi_model.py

@@ -4,11 +4,12 @@ import torch.nn as nn
 import torch.nn.functional as F
 import tqdm
 from transformers import LlamaModel, LlamaConfig
+from transformers.modeling_utils import ModuleUtilsMixin
 
 from midi_tokenizer import MIDITokenizer
 
 
-class MIDIModel(nn.Module):
+class MIDIModel(nn.Module, ModuleUtilsMixin):
     def __init__(self, tokenizer: MIDITokenizer, n_layer=12, n_head=16, n_embd=1024, n_inner=4096, flash=False,
                  *args, **kwargs):
         super(MIDIModel, self).__init__()