chore: refactor _local_predict for better performances (#171)

Signed-off-by: Mathieu Andreux <[email protected]>
Substra · Sep 21, 2023 · 16d6c85 · 16d6c85
1 parent 5ee6f3a
commit 16d6c85
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/substrafl/algorithms/pytorch/torch_base_algo.py b/substrafl/algorithms/pytorch/torch_base_algo.py
@@ -158,11 +158,9 @@ def _local_predict(self, predict_dataset: torch.utils.data.Dataset, predictions_
 
         self._model.eval()
 
-        predictions = torch.Tensor([]).to(self._device)
         with torch.inference_mode():
-            for x in predict_loader:
-                x = x.to(self._device)
-                predictions = torch.cat((predictions, self._model(x)), 0)
+            predictions = [self._model(x.to(self._device)) for x in predict_loader]
+        predictions = torch.cat(predictions, dim=0)
 
         predictions = predictions.cpu().detach()
         self._save_predictions(predictions, predictions_path)