keep best models on cpu

spozdn · Jan 13, 2024 · ef223d1 · ef223d1
1 parent 2eb5dfa
commit ef223d1
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/utilities.py b/src/utilities.py
@@ -52,7 +52,10 @@ def __init__(self):
     def update(self, model_now, error_now, epoch_now, additional_info=None):
         if (self.best_error is None) or (error_now < self.best_error):
             self.best_error = error_now
+            model_now.to('cpu')
             self.best_model = copy.deepcopy(model_now)
+            device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+            model_now.to(device)
             self.best_epoch = epoch_now
             self.additional_info = additional_info