clear cache when valdiating

OuyangWenyu · Jan 3, 2025 · 01c849b · 01c849b
1 parent ae12fff
commit 01c849b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/torchhydro/trainers/train_utils.py b/torchhydro/trainers/train_utils.py
@@ -1,10 +1,10 @@
 """
 Author: Wenyu Ouyang
 Date: 2024-04-08 18:16:26
-LastEditTime: 2025-01-01 15:53:15
+LastEditTime: 2025-01-03 20:32:42
 LastEditors: Wenyu Ouyang
 Description: Some basic functions for training
-FilePath: \torchhydro\torchhydro\trainers\train_utils.py
+FilePath: /torchhydro/torchhydro/trainers/train_utils.py
 Copyright (c) 2024-2024 Wenyu Ouyang. All rights reserved.
 """
 
@@ -476,6 +476,8 @@ def compute_validation(
             trg, output = model_infer(seq_first, device, model, src, trg)
             obs.append(trg)
             preds.append(output)
+            # clear memory to save GPU memory
+            torch.cuda.empty_cache()
         # first dim is batch
         obs_final = torch.cat(obs, dim=0)
         pred_final = torch.cat(preds, dim=0)