add clear gpu cache after each batch in eval

OuyangWenyu · Dec 14, 2024 · 2b9c0d3 · 2b9c0d3
1 parent 184534d
commit 2b9c0d3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torchhydro/trainers/deep_hydro.py b/torchhydro/trainers/deep_hydro.py
@@ -386,6 +386,8 @@ def inference(self) -> Tuple[xr.Dataset, xr.Dataset]:
                 ys, pred = model_infer(seq_first, device, self.model, xs, ys)
                 test_preds.append(pred.cpu().numpy())
                 obss.append(ys.cpu().numpy())
+                # clear GPU cache -- this will not impact results
+                torch.cuda.empty_cache()
             pred = reduce(lambda x, y: np.vstack((x, y)), test_preds)
             obs = reduce(lambda x, y: np.vstack((x, y)), obss)
         if pred.ndim == 2: