added cpu parameter support (#423)

2022-08-05 22:10:21 +01:00
parent 72a9c9628a
commit 1560034c77
1 changed files with 1 additions and 1 deletions
@@ -65,7 +65,7 @@ def custom(path_or_model='path/to/model.pt', autoshape=True):
    Returns:
        pytorch model
    """
-    model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model  # load checkpoint
+    model = torch.load(path_or_model, map_location=torch.device('cpu')) if isinstance(path_or_model, str) else path_or_model  # load checkpoint
    if isinstance(model, dict):
        model = model['ema' if model.get('ema') else 'model']  # load model