From f7c887703ae1d35da4a2d349b57ee88bbc11b7ea Mon Sep 17 00:00:00 2001 From: Danielle Pintz Date: Mon, 6 Nov 2023 18:53:29 -0800 Subject: [PATCH] Remove __enter__ and __exit__ from MemorySnapshotProfiler (#611) Summary: There is no need for MemorySnapshotProfiler to be a context manager since it conflicts with start_step and stop_step Differential Revision: D51049497 --- torchtnt/utils/memory_snapshot_profiler.py | 14 +------------- 1 file changed, 1 insertion(+), 13 deletions(-) diff --git a/torchtnt/utils/memory_snapshot_profiler.py b/torchtnt/utils/memory_snapshot_profiler.py index 58f2b69b99..07c31f0c77 100644 --- a/torchtnt/utils/memory_snapshot_profiler.py +++ b/torchtnt/utils/memory_snapshot_profiler.py @@ -6,8 +6,7 @@ import logging from dataclasses import dataclass -from types import TracebackType -from typing import Optional, Type +from typing import Optional import torch from torchtnt.utils.oom import attach_oom_observer, log_memory_snapshot @@ -115,17 +114,6 @@ def __init__( f"Created MemorySnapshotProfiler with MemorySnapshotParams={self.params}." ) - def __enter__(self) -> None: - self.start() - - def __exit__( - self, - exc_type: Optional[Type[BaseException]], - exc_value: Optional[BaseException], - tb: Optional[TracebackType], - ) -> Optional[bool]: - self.stop() - def start(self) -> None: if not torch.cuda.is_available(): logger.warn("CUDA unavailable. Not recording memory history.")