zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

.to(device) bug fix

This commit is contained in:
Glenn Jocher 2020-07-05 23:37:36 -07:00
parent a62333e808
commit 6b95d6d2c0
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
test.py
View File

@ -31,7 +31,7 @@ def test(data,
# Load model
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
# Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99