File tree Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Expand file tree Collapse file tree 1 file changed +2
-2
lines changed Original file line number Diff line number Diff line change 77
88def load_checkpoint (model , checkpoint_path , use_ema = False ):
99 if checkpoint_path and os .path .isfile (checkpoint_path ):
10- checkpoint = torch .load (checkpoint_path )
10+ checkpoint = torch .load (checkpoint_path , map_location = 'cpu' )
1111 state_dict_key = ''
1212 if isinstance (checkpoint , dict ):
1313 state_dict_key = 'state_dict'
@@ -32,7 +32,7 @@ def resume_checkpoint(model, checkpoint_path):
3232 optimizer_state = None
3333 resume_epoch = None
3434 if os .path .isfile (checkpoint_path ):
35- checkpoint = torch .load (checkpoint_path )
35+ checkpoint = torch .load (checkpoint_path , map_location = 'cpu' )
3636 if isinstance (checkpoint , dict ) and 'state_dict' in checkpoint :
3737 new_state_dict = OrderedDict ()
3838 for k , v in checkpoint ['state_dict' ].items ():
You can’t perform that action at this time.
0 commit comments