Use binary units for CUDA memory summary

milesial · pytorchmergebot · commit 686d2ce167a6 · 2023-01-14T01:32:11.000Z
diff --git a/torch/cuda/memory.py b/torch/cuda/memory.py
@@ -453,15 +453,15 @@ def memory_summary(device: Union[Device, int] = None, abbreviated: bool = False)
     stats = memory_stats(device=device)
 
     def _format_size(sz, pref_sz):
-        prefixes = ["B ", "KB", "MB", "GB", "TB", "PB"]
+        prefixes = ["B  ", "KiB", "MiB", "GiB", "TiB", "PiB"]
         prefix = prefixes[0]
         for new_prefix in prefixes[1:]:
             if pref_sz < 768 * 1024:
                 break
             prefix = new_prefix
             sz //= 1024
             pref_sz /= 1024
-        return "{:7d} {}".format(sz, prefix)
+        return "{:6d} {}".format(sz, prefix)
 
     def _format_count(cnt, pref_cnt):
         prefixes = [" ", "K", "M"]
