Commit
·
185b654
1
Parent(s):
6b2c7e9
add register
Browse files- custom_model.py +3 -0
custom_model.py
CHANGED
|
@@ -18,3 +18,6 @@ class CustomModel(PreTrainedModel):
|
|
| 18 |
def forward(self, input_ids):
|
| 19 |
output = self.linear(input_ids)
|
| 20 |
return output
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
def forward(self, input_ids):
|
| 19 |
output = self.linear(input_ids)
|
| 20 |
return output
|
| 21 |
+
|
| 22 |
+
AutoConfig.register("custom-model", CustomModelConfig)
|
| 23 |
+
AutoModel.register(CustomModelConfig, CustomModel)
|