torch-memory-saver 0.0.2__tar.gz → 0.0.4__tar.gz

{torch_memory_saver-0.0.2/torch_memory_saver.egg-info → torch_memory_saver-0.0.4}/PKG-INFO

@@ -1,5 +1,5 @@
 Metadata-Version: 2.1
 Name: torch_memory_saver
-Version: 0.0.2
+Version: 0.0.4
 Requires-Python: >=3.9
 License-File: LICENSE

{torch_memory_saver-0.0.2 → torch_memory_saver-0.0.4}/setup.py

@@ -7,7 +7,7 @@ logger = logging.getLogger(__name__)
 
 setup(
     name='torch_memory_saver',
-    version='0.0.2',
+    version='0.0.4',
     ext_modules=[setuptools.Extension(
         'torch_memory_saver_cpp',
         ['csrc/torch_memory_saver.cpp'],

{torch_memory_saver-0.0.2 → torch_memory_saver-0.0.4}/torch_memory_saver/__init__.py

@@ -13,13 +13,14 @@ logger = logging.getLogger(__name__)
 
 class TorchMemorySaver:
     def __init__(self):
-        self._mem_pool = torch.cuda.MemPool()
+        self._mem_pool = None
         self._id = _global_info.next_id()
         assert self._id == 1, 'Only support one single instance yet (multi-instance will be implemented later)'
 
     @contextmanager
     def region(self):
         if _global_info.binary_info.enabled:
+            self._ensure_mem_pool()
             with torch.cuda.use_mem_pool(self._mem_pool):
                 _global_info.binary_info.cdll.tms_region_enter()
                 try:
@@ -37,6 +38,14 @@ class TorchMemorySaver:
         if _global_info.binary_info.enabled:
             _global_info.binary_info.cdll.tms_resume()
 
+    @property
+    def enabled(self):
+        return _global_info.binary_info.enabled
+
+    def _ensure_mem_pool(self):
+        if self._mem_pool is None:
+            self._mem_pool = torch.cuda.MemPool()
+
 
 @dataclass
 class _BinaryInfo:

{torch_memory_saver-0.0.2 → torch_memory_saver-0.0.4/torch_memory_saver.egg-info}/PKG-INFO

@@ -1,5 +1,5 @@
 Metadata-Version: 2.1
 Name: torch_memory_saver
-Version: 0.0.2
+Version: 0.0.4
 Requires-Python: >=3.9
 License-File: LICENSE