.half() bug fix

This commit is contained in:
Glenn Jocher 2020-04-22 10:46:26 -07:00
parent 03c6a2d6fa
commit 2f636d5740
1 changed files with 2 additions and 1 deletions

View File

@ -75,7 +75,8 @@ def detect(save_img=False):
# Run inference
t0 = time.time()
_ = model(torch.zeros((1, 3, img_size, img_size), device=device)) if device.type != 'cpu' else None # run once
img = torch.zeros((1, 3, img_size, img_size), device=device) # init img
_ = model(img.half() if half else img.float()) if device.type != 'cpu' else None # run once
for path, img, im0s, vid_cap in dataset:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32