fix(camelyon/weldon_fedavg): can't convert cuda:0 device type tensor …

…to numpy (#245) Signed-off-by: Thibault Camalon <[email protected]>
Substra · Sep 4, 2024 · 7a83422 · 7a83422
1 parent 92ddbf3
commit 7a83422
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/benchmark/camelyon/weldon_fedavg.py b/benchmark/camelyon/weldon_fedavg.py
@@ -89,7 +89,7 @@ def _local_predict(self, predict_dataset):
                 for X in dataloader:
                     y_pred.append(self._model(X))
 
-            y_pred = torch.cat(y_pred).numpy()
+            y_pred = torch.cat(y_pred).cpu().numpy()
 
             return y_pred
 

diff --git a/changes/245.fixed b/changes/245.fixed
@@ -0,0 +1 @@
+Use Tensor.cpu() to copy the tensor to host memory first in Camelyon benchmark
Original file line number	Diff line number	Diff line change
		@@ -0,0 +1 @@
		Use Tensor.cpu() to copy the tensor to host memory first in Camelyon benchmark