Browse Source

`attempt_load()` deserialize fix (#8051)

modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
258f4c6438
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions
  1. +2
    -2
      models/experimental.py

+ 2
- 2
models/experimental.py View File

@@ -77,8 +77,8 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
ckpt = torch.load(attempt_download(w))
ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float() # FP32 model
ckpt = torch.load(attempt_download(w), map_location=device)
ckpt = (ckpt.get('ema') or ckpt['model']).float() # FP32 model
model.append(ckpt.fuse().eval() if fuse else ckpt.eval()) # fused or un-fused model in eval mode

# Compatibility updates

Loading…
Cancel
Save