Improve GPU name (#5478)

* Improve GPU name * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py
2021-11-03 16:59:53 +01:00 · 2021-11-03 16:59:53 +01:00 · df30426c03
parent 0eb37ad8af
commit df30426c03
1 changed files with 1 additions and 1 deletions
--- 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'