提交 d1e57166 authored 作者: Glenn Jocher's avatar Glenn Jocher

train with multi-gpu half test bug fix #99

上级 1dfc2852
...@@ -46,7 +46,7 @@ def test(data, ...@@ -46,7 +46,7 @@ def test(data,
else: # called by train.py else: # called by train.py
training = True training = True
device = next(model.parameters()).device # get model device device = next(model.parameters()).device # get model device
half = device.type != 'cpu' # half precision only supported on CUDA half = device.type != 'cpu' and torch.cuda.device_count() == 1 # half precision only supported on single-GPU
if half: if half:
model.half() # to FP16 model.half() # to FP16
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论