Update train.py

janelia-cellmap · Feb 7, 2024 · 3c5f2da · 3c5f2da
1 parent 1490440
commit 3c5f2da
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/dacapo/train.py b/dacapo/train.py
@@ -159,7 +159,7 @@ def train_run(
 
             run.model.eval()
             # free up optimizer memory to allow larger validation blocks
-            # run.model = run.model.to(torch.device("cpu"))
+            run.model = run.model.to(torch.device("cpu"))
             run.move_optimizer(torch.device("cpu"), empty_cuda_cache=True)
 
             weights_store.store_weights(run, iteration_stats.iteration + 1)