Unverified 提交 df30426c authored 作者: Glenn Jocher's avatar Glenn Jocher 提交者: GitHub

Improve GPU name (#5478)

* Improve GPU name * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py * Update torch_utils.py
上级 0eb37ad8
...@@ -75,7 +75,7 @@ def select_device(device='', batch_size=None): ...@@ -75,7 +75,7 @@ def select_device(device='', batch_size=None):
space = ' ' * (len(s) + 1) space = ' ' * (len(s) + 1)
for i, d in enumerate(devices): for i, d in enumerate(devices):
p = torch.cuda.get_device_properties(i) 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: else:
s += 'CPU\n' s += 'CPU\n'
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论