charlesxsh commited on
Commit
55fcde9
·
1 Parent(s): 3bf08d4
Files changed (1) hide show
  1. custom_model.py +1 -0
custom_model.py CHANGED
@@ -9,6 +9,7 @@ class CustomModelConfig(PretrainedConfig):
9
  super().__init__(**kwargs)
10
  self.hidden_size = hidden_size
11
  self.num_attention_heads = 1
 
12
 
13
 
14
  class CustomModel(PreTrainedModel):
 
9
  super().__init__(**kwargs)
10
  self.hidden_size = hidden_size
11
  self.num_attention_heads = 1
12
+ self.vocab_size = 1000
13
 
14
 
15
  class CustomModel(PreTrainedModel):