Update modeling_fake_custom.py

#2
by Rocketknight1 HF staff - opened
Files changed (1) hide show
  1. modeling_fake_custom.py +2 -0
modeling_fake_custom.py CHANGED
@@ -2,8 +2,10 @@ from transformers import GPT2Model, GPT2LMHeadModel
2
  from .configuration_fake_custom import FakeCustomConfig
3
 
4
  class FakeCustomModel(GPT2Model):
 
5
  config_class = FakeCustomConfig
6
 
7
 
8
  class FakeCustomModelForCausalLM(GPT2LMHeadModel):
 
9
  config_class = FakeCustomConfig
 
2
  from .configuration_fake_custom import FakeCustomConfig
3
 
4
  class FakeCustomModel(GPT2Model):
5
+ _auto_class = "AutoModel"
6
  config_class = FakeCustomConfig
7
 
8
 
9
  class FakeCustomModelForCausalLM(GPT2LMHeadModel):
10
+ _auto_class = "AutoModelForCausalLM"
11
  config_class = FakeCustomConfig