prediction fix

GilesStrong · Nov 30, 2020 · 9a45395 · 9a45395
1 parent c5e931f
commit 9a45395
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/CHANGES.md b/CHANGES.md
@@ -13,6 +13,7 @@
 ## Fixes
 
 - `Model` now creates `cb_savepath` is it didn't already exist
+- Bug in `PredHandler` where predictions were kept on device leading to increased memory usage
 
 ## Changes
 

diff --git a/lumin/nn/callbacks/pred_handlers.py b/lumin/nn/callbacks/pred_handlers.py
@@ -14,4 +14,4 @@ def on_pred_begin(self) -> None: self.preds = []
     def on_pred_end(self) -> None: self.preds = torch.cat(self.preds)
     def get_preds(self) -> np.ndarray: return self.preds
     def on_forwards_end(self) -> None:
-        if self.model.fit_params.state == 'test': self.preds.append(self.model.fit_params.y_pred)
+        if self.model.fit_params.state == 'test': self.preds.append(self.model.fit_params.y_pred.cpu().detach())