Add free_memory hook for pytorch

Octomiro · Jan 24, 2024 · fcc1fce · fcc1fce
1 parent 8d77aca
commit fcc1fce
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git a/octopipes/hooks/__init__.py b/octopipes/hooks/__init__.py
diff --git a/octopipes/hooks/pytorch_hooks.py b/octopipes/hooks/pytorch_hooks.py
@@ -0,0 +1,10 @@
+import gc
+
+import torch
+
+
+def free_memory(wf_iter):
+    del wf_iter
+    torch.cuda.empty_cache()
+    gc.collect()
+
diff --git a/setup.py b/setup.py
@@ -26,7 +26,8 @@
         'cmap'
     ],
     extras_require={
-        'opencv': ['opencv-python']
+        'opencv': ['opencv-python'],
+        'pytorch': ['torch', 'torchvision']
     },
     classifiers=[
         'Programming Language :: Python :: 3.11',