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
|
||||
training = True
|
||||
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:
|
||||
model.half() # to FP16
|
||||
|
||||
|
||||
Reference in new issue
Block a user