fix: tensor conversion with gpu

The conversion of GPU tensors to numpy arrays caused an error, which should now be fixed.
online-ml · Aug 30, 2023 · 84fca25 · 84fca25
1 parent 50a240d
commit 84fca25
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/deep_river/utils/tensor_conversion.py b/deep_river/utils/tensor_conversion.py
@@ -156,7 +156,7 @@ def output2proba(
         else:
             preds = torch.sigmoid(preds)
 
-    preds_np = preds.detach().numpy()
+    preds_np = preds.numpy(force=True)
     if preds_np.shape[1] == 1:
         preds_np = np.hstack((preds_np, 1 - preds_np))
     n_unobserved_classes = preds_np.shape[1] - len(classes)