Clarify units are MiB.

rapidsai · Sep 6, 2024 · 8075ba3 · 8075ba3
1 parent 1b83eb1
commit 8075ba3
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/cucim/tests/performance/clara/test_read_region_memory_usage.py b/python/cucim/tests/performance/clara/test_read_region_memory_usage.py
@@ -23,7 +23,8 @@
 
 
 def test_read_region_cuda_memleak(testimg_tiff_stripe_4096x4096_256_jpeg):
-    def get_used_gpu_memory():
+    def get_used_gpu_memory_mib():
+        """Get the used GPU memory in MiB."""
         status, free, total = cuda.cudart.cudaMemGetInfo()
         if status != cuda.cudart.cudaError_t.cudaSuccess:
             raise RuntimeError("Failed to get GPU memory info.")
@@ -36,11 +37,11 @@ def get_used_gpu_memory():
 
     img = open_image_cucim(testimg_tiff_stripe_4096x4096_256_jpeg)
 
-    mem_usage_history = [get_used_gpu_memory()]
+    mem_usage_history = [get_used_gpu_memory_mib()]
 
     for i in range(10):
         _ = img.read_region(device="cuda")
-        mem_usage_history.append(get_used_gpu_memory())
+        mem_usage_history.append(get_used_gpu_memory_mib())
 
     print(mem_usage_history)