提交 6b95d6d2 authored 作者: Glenn Jocher's avatar Glenn Jocher

.to(device) bug fix

上级 a62333e8
......@@ -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
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论