gsar78 commited on
Commit
66ee760
1 Parent(s): b4c4b45

Update custom_model_package/custom_model.py

Browse files
custom_model_package/custom_model.py CHANGED
@@ -10,7 +10,7 @@ class CustomConfig(PretrainedConfig):
10
  super().__init__(**kwargs)
11
  self.num_emotion_labels = num_emotion_labels
12
 
13
- class CustomModel(nn.Module, PyTorchModelHubMixin):
14
  config_class = CustomConfig
15
 
16
  def __init__(self, config, num_emotion_labels):
 
10
  super().__init__(**kwargs)
11
  self.num_emotion_labels = num_emotion_labels
12
 
13
+ class CustomModel(XLMRobertaForSequenceClassification):
14
  config_class = CustomConfig
15
 
16
  def __init__(self, config, num_emotion_labels):