Update custom_model_package/custom_model.py

custom_model_package/custom_model.py

@@ -1,6 +1,7 @@
 from transformers import PretrainedConfig, XLMRobertaForSequenceClassification
 import torch.nn as nn
 import torch
+from huggingface_hub import PyTorchModelHubMixin
 
 class CustomConfig(PretrainedConfig):
     model_type = "custom_model"
@@ -9,7 +10,7 @@ class CustomConfig(PretrainedConfig):
         super().__init__(**kwargs)
         self.num_emotion_labels = num_emotion_labels
 
-class CustomModel(PreTrainedModel):
+class CustomModel(nn.Module, PyTorchModelHubMixin):
     config_class = CustomConfig
     
     def __init__(self, config, num_emotion_labels):