Hector Lopez commited on
Commit
3bf6581
1 Parent(s): 3fa54be

refactor: Added map_location parameter into get_models()

Browse files
Files changed (1) hide show
  1. utils.py +2 -1
utils.py CHANGED
@@ -73,7 +73,8 @@ def get_models(
73
  backbone_name='d0',
74
  img_size=512,
75
  classes=['Waste'],
76
- revise_keys=[(r'^model\.', '')])
 
77
 
78
  det_model = checkpoint_and_model['model']
79
  det_model.eval()
 
73
  backbone_name='d0',
74
  img_size=512,
75
  classes=['Waste'],
76
+ revise_keys=[(r'^model\.', '')],
77
+ map_location='cpu')
78
 
79
  det_model = checkpoint_and_model['model']
80
  det_model.eval()