Fix node classification loss

THUDM · Aug 4, 2020 · c876d9e · c876d9e
1 parent 6e33d80
commit c876d9e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cogdl/tasks/node_classification.py b/cogdl/tasks/node_classification.py
@@ -91,7 +91,7 @@ def _test_step(self, split="val"):
             mask = self.data.val_mask
         else:
             mask = self.data.test_mask
-        loss = F.nll_loss(logits[mask], self.data.y[mask])
+        loss = F.nll_loss(logits[mask], self.data.y[mask]).item()
 
         pred = logits[mask].max(1)[1]
         acc = pred.eq(self.data.y[mask]).sum().item() / mask.sum().item()