yo

config.json

@@ -9,5 +9,5 @@
   "architectures": [
     "CustomLM"
   ],
-  "num_attention_heads": 4
+  "num_attention_heads": 1
 }

custom_model.py

@@ -8,6 +8,7 @@ class CustomModelConfig(PretrainedConfig):
     def __init__(self, hidden_size=128, **kwargs):
         super().__init__(**kwargs)
         self.hidden_size = hidden_size
+        self.num_attention_heads = 1
 
 
 class CustomModel(PreTrainedModel):