train with multi-gpu half test bug fix #99
This commit is contained in:
1 parent
1dfc28527f
commit
d1e5716610
1 file changed
+1
-1
@@ -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
|
||||||
|
|
||||||
|
|||||||
Reference in new issue
Block a user