Update modeling_me2bert.py

modeling_me2bert.py

@@ -187,11 +187,11 @@ class ME2BertModel(PreTrainedModel):
 
         if emotion_features is not None:
             emotion_features = emotion_features[:gated_output.shape[0], :]
-            class_output = torch.cat((gated_output, domain_feature, emotion_features), dim=1)
+            class_output = torch.cat((gated_output.to(device), domain_feature.to(device), emotion_features.to(device)), dim=1)
 
         else:
             emotion_features = torch.zeros(gated_output.shape[0], self.emotion_dim).to(device)
-            class_output = torch.cat((gated_output, domain_feature, emotion_features), dim=1)
+            class_output = torch.cat((gated_output.to(device), domain_feature.to(device), emotion_features.to(device)), dim=1)
 
         class_output = torch.sigmoid(self.mf_classifier(class_output))
         if return_dict: