extra helper

drisspg · Jan 17, 2024 · 61830d7 · 61830d7
1 parent ba90b6a
commit 61830d7
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/transformer_nuggets/utils/benchmark.py b/transformer_nuggets/utils/benchmark.py
@@ -79,6 +79,14 @@ def profile_function(
     return prof
 
 
+@contextmanager
+def print_max_memory_usage():
+    try:
+        yield
+    finally:
+        print(f"Max Cuda Memory Used: {torch.cuda.max_memory_allocated() / (1024**3):.4f} GiB")
+
+
 @contextmanager
 def print_cuda_memory_usage():
     initial_memory = torch.cuda.memory_allocated()