glenn-jocher commited on
Commit
6b95d6d
·
1 Parent(s): a62333e

.to(device) bug fix

Browse files
Files changed (1) hide show
  1. test.py +1 -1
test.py CHANGED
@@ -31,7 +31,7 @@ def test(data,
31
 
32
  # Load model
33
  google_utils.attempt_download(weights)
34
- model = torch.load(weights, map_location=device)['model'].float().fuse() # load to FP32
35
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
36
 
37
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
 
31
 
32
  # Load model
33
  google_utils.attempt_download(weights)
34
+ model = torch.load(weights, map_location=device)['model'].float().fuse().to(device) # load to FP32
35
  imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
36
 
37
  # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99