Update custom_model_package/custom_model.py

custom_model_package/custom_model.py

@@ -56,8 +56,7 @@ class CustomModel(XLMRobertaForSequenceClassification):
         if config is None:
             config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
         num_emotion_labels = kwargs.pop('num_emotion_labels', config.num_emotion_labels)
-        model = super(CustomModel, cls).from_pretrained(pretrained_model_name_or_path, config=config, *model_args, **kwargs)
-        model.num_emotion_labels = num_emotion_labels
+        model = super(CustomModel, cls).from_pretrained(pretrained_model_name_or_path, config=config, num_emotion_labels=num_emotion_labels, *model_args, **kwargs)
         return model
 
 # Register the custom configuration and model