mem clear (#440)

* mem clear * mem clear
2026-09-09 11:32:07 +08:00 · 2019-10-30 12:11:21 -04:00
parent 4df4d4cc03
commit 8347a6c87e
1 changed files with 7 additions and 0 deletions
@@ -32,10 +32,17 @@ class TrainerIOMixin(object):
        :param model:
        :return:
        """
+        # clear cache before restore
+        if self.on_gpu:
+            torch.cuda.empty_cache()

        # if script called from hpc resubmit, load weights
        did_restore_hpc_weights = self.restore_hpc_weights_if_needed(model)

+        # clear cache after restore
+        if self.on_gpu:
+            torch.cuda.empty_cache()
+
        if not did_restore_hpc_weights:
            # restore weights if same exp version
            self.restore_state_if_checkpoint_exists(model)