Skip to content

Commit 1560034

Browse files
authored
added cpu parameter support (#423)
1 parent 72a9c96 commit 1560034

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

hubconf.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ def custom(path_or_model='path/to/model.pt', autoshape=True):
6565
Returns:
6666
pytorch model
6767
"""
68-
model = torch.load(path_or_model) if isinstance(path_or_model, str) else path_or_model # load checkpoint
68+
model = torch.load(path_or_model, map_location=torch.device('cpu')) if isinstance(path_or_model, str) else path_or_model # load checkpoint
6969
if isinstance(model, dict):
7070
model = model['ema' if model.get('ema') else 'model'] # load model
7171

0 commit comments

Comments
 (0)