Update modeling_fake_custom.py

modeling_fake_custom.py

@@ -1,9 +1,13 @@
-from transformers import GPT2Model, GPT2LMHeadModel
+from transformers import PretrainedModel
 from .configuration_fake_custom import FakeCustomConfig
 
-class FakeCustomModel(GPT2Model):
+class FakeCustomPreTrainedModel(PreTrainedModel):
     config_class = FakeCustomConfig
 
 
-class FakeCustomModelForCausalLM(GPT2LMHeadModel):
-    config_class = FakeCustomConfig
+class FakeCustomModel(FakeCustomPreTrainedModel):
+    _auto_class = "AutoModel"
+
+
+class FakeCustomForCausalLM(FakeCustomPreTrainedModel):
+    _auto_class = "AutoModelForCausalLM"