Update model name in auto_map

backpack_model.py

@@ -148,7 +148,7 @@ class BackpackGPT2Model(BackpackGPT2PreTrainedModel):
     def get_sense_network(self):
         return self.sense_network
 
-    def forward(self, input_ids, position_ids):
+    def forward(self, input_ids, position_ids: Optional[torch.LongTensor] = None):
         # Compute senses
         sense_input_embeds = self.word_embeddings(input_ids)
         senses = self.sense_network(sense_input_embeds) # (bs, nv, s, d)

config.json

@@ -76,7 +76,7 @@
     "sense_intermediate_scale": 4,
     "auto_map": {
         "AutoConfig": "backpack_config.BackpackGPT2Config",
-        "AutoModelForCausalLM": "backpack_model.BackpackGPT2Model"
+        "AutoModelForCausalLM": "backpack_model.BackpackGPT2LMHeadModel"
     },
     "model_type": "backpack-gpt2"
 }