diff --git a/utils/torch_utils.py b/utils/torch_utils.py index fc21414..ed399a7 100644 --- 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'