empty cuda cache after epoch instead of batch

batra-mlp-lab · May 28, 2020 · 0e010a4 · 0e010a4
1 parent 6c4e2b8
commit 0e010a4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/train.py b/train.py
@@ -280,7 +280,7 @@ def lr_lambda_fun(current_iteration: int) -> float:
 
         scheduler.step(global_iteration_step)
         global_iteration_step += 1
-        torch.cuda.empty_cache()
+    torch.cuda.empty_cache()
 
     # -------------------------------------------------------------------------
     #   ON EPOCH END  (checkpointing and validation)