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

Update torch_utils.py (#1895)

上级 0f11aaf5
...@@ -61,7 +61,7 @@ def select_device(device='', batch_size=None): ...@@ -61,7 +61,7 @@ def select_device(device='', batch_size=None):
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
cuda = torch.cuda.is_available() and not cpu cuda = not cpu and torch.cuda.is_available()
if cuda: if cuda:
n = torch.cuda.device_count() n = torch.cuda.device_count()
if n > 1 and batch_size: # check that batch_size is compatible with device_count if n > 1 and batch_size: # check that batch_size is compatible with device_count
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论