Update

Author: Yan98

train.py

@@ -168,7 +168,7 @@ def training_step(self,data,idx):
 
         pred_count = self.model(data.x_dict,data.edge_index_dict, emb, size)
         loss   = self.criterion(pred_count,data["window"]["y"][:,mask])
-        corrloss = self.correlationMetric(pred_count,data["window"]["y"])
+        corrloss = self.correlationMetric(pred_count,data["window"]["y"][:,mask])
 
         optimizer.zero_grad()
         self.manual_backward(loss + corrloss * 0.5)