소스 검색

.to(device) bug fix

5.0
Glenn Jocher 4 년 전
부모
커밋
6b95d6d2c0
1개의 변경된 파일1개의 추가작업 그리고 1개의 파일을 삭제
  1. +1
    -1
      test.py

+ 1
- 1
test.py 파일 보기

@@ -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

Loading…
취소
저장