return tensor instead

kumo-ai · Oct 21, 2024 · d506b54 · d506b54
1 parent 97bacbd
commit d506b54
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/examples/ijcai_example.py b/examples/ijcai_example.py
@@ -334,7 +334,7 @@ def test(loader: NeighborLoader, desc: str) -> np.ndarray:
 
         _, pred_mini = torch.topk(scores, k=args.eval_k, dim=1)
         pred_list.append(pred_mini)
-    pred = torch.cat(pred_list, dim=0).cpu().numpy()
+    pred = torch.cat(pred_list, dim=0)
     return pred