FIx downstream impact by the replacement of torch.load in ultralytics package (#13623)
* Update torch_load * Auto-format by https://ultralytics.com/actions --------- Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
1 parent
79c4c31d16
commit
2540fd4c1c
9 files changed
+20
-12
No files matched your search
@@ -6,6 +6,7 @@ import math
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from ultralytics.utils.patches import torch_load
|
||||
|
||||
from utils.downloads import attempt_download
|
||||
|
||||
@@ -95,7 +96,7 @@ def attempt_load(weights, device=None, inplace=True, fuse=True):
|
||||
|
||||
model = Ensemble()
|
||||
for w in weights if isinstance(weights, list) else [weights]:
|
||||
ckpt = torch.load(attempt_download(w), map_location="cpu") # load
|
||||
ckpt = torch_load(attempt_download(w), map_location="cpu") # load
|
||||
ckpt = (ckpt.get("ema") or ckpt["model"]).to(device).float() # FP32 model
|
||||
|
||||
# Model compatibility updates
|
||||
|
||||
Reference in new issue
Block a user