Browse Source

FP16 to FP32 ckpt load

5.0
Glenn Jocher 4 years ago
parent
commit
76ca367a01
2 changed files with 2 additions and 2 deletions
  1. +1
    -1
      detect.py
  2. +1
    -1
      test.py

+ 1
- 1
detect.py View File

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

+ 1
- 1
test.py View File

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

Loading…
Cancel
Save