Spaces:

xiang-wuu
/

yolov5

Runtime error

glenn-jocher commited on Jun 16, 2020

Commit

d1e5716

•

1 Parent(s): 1dfc285

train with multi-gpu half test bug fix #99

Files changed (1) hide show

test.py CHANGED Viewed

@@ -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
         if half:
             model.half()  # to FP16

     else:  # called by train.py
         training = True
         device = next(model.parameters()).device  # get model device
+        half = device.type != 'cpu' and torch.cuda.device_count() == 1  # half precision only supported on single-GPU
         if half:
             model.half()  # to FP16