diff --git a/detect.py b/detect.py index 86b6169..bf858e4 100644 --- a/detect.py +++ b/detect.py @@ -18,7 +18,7 @@ def detect(save_img=False): # Load model google_utils.attempt_download(weights) - model = torch.load(weights, map_location=device)['model'] + model = torch.load(weights, map_location=device)['model'].float() # load to FP32 # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning # model.fuse() model.to(device).eval() diff --git a/test.py b/test.py index 292c4c4..60faf23 100644 --- a/test.py +++ b/test.py @@ -32,7 +32,7 @@ def test(data, # Load model google_utils.attempt_download(weights) - model = torch.load(weights, map_location=device)['model'] + model = torch.load(weights, map_location=device)['model'].float() # load to FP32 torch_utils.model_info(model) # model.fuse() model.to(device)