Fix scoring function

PreferredAI · Jan 17, 2024 · 78b8fd2 · 78b8fd2
1 parent 784841f
commit 78b8fd2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/cornac/models/dnntsp/dnntsp.py b/cornac/models/dnntsp/dnntsp.py
@@ -636,4 +636,4 @@ def score(model: TemporalSetPrediction, history_baskets, total_items, device="cp
         is_test=True,
     )
     preds = model(g, nodes_feature, edges_weight, lengths, nodes)
-    return preds[0].cpu().detach().numpy()
+    return preds.cpu().detach().numpy()