cpu to device select

2025-01-13 18:11:56 +08:00
parent a50f777839
commit 744fb7b7b2
4 changed files with 3 additions and 3 deletions
--- a/models/pycache/experimental.cpython-39.pyc
+++ b/models/pycache/experimental.cpython-39.pyc
--- a/models/experimental.py
+++ b/models/experimental.py
@ -76,7 +76,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=device)  # load
        ckpt = (ckpt.get('ema') or ckpt['model']).to(device).float()  # FP32 model

        # Model compatibility updates