Improve GPU name (ultralytics#5478)

* Improve GPU name * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py
fhkiel-mlaip · Nov 3, 2021 · 238cccf · 238cccf
1 parent ee20ce8
commit 238cccf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
@@ -75,7 +75,7 @@ def select_device(device='', batch_size=None):
         space = ' ' * (len(s) + 1)
         for i, d in enumerate(devices):
             p = torch.cuda.get_device_properties(i)
-            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n"  # bytes to MB
+            s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2:.0f}MiB)\n"  # bytes to MB
     else:
         s += 'CPU\n'