use_cache=False to avoid warning

SkyTNT · Oct 5, 2024 · 673d601 · 673d601
1 parent 58fb183
commit 673d601
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/midi_model.py b/midi_model.py
@@ -29,11 +29,13 @@ def get_config(tokenizer_ver="v2", optimise_midi=True, n_layer=12, n_head=16, n_
         net_config = LlamaConfig(vocab_size=tokenizer.vocab_size,
                                  hidden_size=n_embd, num_attention_heads=n_head,
                                  num_hidden_layers=n_layer, intermediate_size=n_inner,
-                                 pad_token_id=tokenizer.pad_id, max_position_embeddings=4096)
+                                 pad_token_id=tokenizer.pad_id, max_position_embeddings=4096,
+                                 use_cache=False)
         net_token_config = LlamaConfig(vocab_size=tokenizer.vocab_size,
                                        hidden_size=n_embd, num_attention_heads=n_head // 4,
                                        num_hidden_layers=n_layer // 4, intermediate_size=n_inner // 4,
-                                       pad_token_id=tokenizer.pad_id, max_position_embeddings=4096)
+                                       pad_token_id=tokenizer.pad_id, max_position_embeddings=4096,
+                                       use_cache=False)
         return MIDIModelConfig(tokenizer, net_config, net_token_config)
 
     @staticmethod