disable recording memory snapshots after dumping

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/611

Disable recording memory snapshots after dumping to files. Otherwise the process won't have a clean shutdown.

Reviewed By: ertrue, wat3rBro

Differential Revision: D48533397

fbshipit-source-id: cbf260823172222b8015008eaffa3d0361fa6233

d2go/utils/gpu_memory_profiler.py

@@ -65,6 +65,8 @@ def log_memory_snapshot(output_dir: str, file_prefix: str = "") -> None:
         logger.info(f"Logging memory snapshot to {save_dir}")
         snapshot = torch.cuda.memory._snapshot()
         dump_snapshot(save_dir, snapshot)
+        # disable recording memory history
+        torch.cuda.memory._record_memory_history(enabled=None)
     except Exception as e:
         logger.error(f"Failed to log memory snapshot to {save_dir}: {e}")