.to(device) bug fix
This commit is contained in:
parent
a62333e808
commit
6b95d6d2c0
2
test.py
2
test.py
|
|
@ -31,7 +31,7 @@ def test(data,
|
||||||
|
|
||||||
# Load model
|
# Load model
|
||||||
google_utils.attempt_download(weights)
|
google_utils.attempt_download(weights)
|
||||||
model = torch.load(weights, map_location=device)['model'].float().fuse() # load to FP32
|
model = torch.load(weights, map_location=device)['model'].float().fuse().to(device) # load to FP32
|
||||||
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
imgsz = check_img_size(imgsz, s=model.stride.max()) # check img_size
|
||||||
|
|
||||||
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue