Nadine Rueegg commited on
Commit
56bd45b
β€’
1 Parent(s): b225901
gradio_demo/barc_demo_v3.py CHANGED
@@ -168,7 +168,7 @@ def run_barc_inference(input_image, bbox=None):
168
  print(path_model_file_complete)
169
  assert os.path.isfile(path_model_file_complete)
170
  print('Loading model weights from file: {}'.format(path_model_file_complete))
171
- checkpoint_complete = torch.load(path_model_file_complete)
172
  state_dict_complete = checkpoint_complete['state_dict']
173
  complete_model.load_state_dict(state_dict_complete, strict=False)
174
  complete_model = complete_model.to(device)
 
168
  print(path_model_file_complete)
169
  assert os.path.isfile(path_model_file_complete)
170
  print('Loading model weights from file: {}'.format(path_model_file_complete))
171
+ checkpoint_complete = torch.load(path_model_file_complete, map_location=device)
172
  state_dict_complete = checkpoint_complete['state_dict']
173
  complete_model.load_state_dict(state_dict_complete, strict=False)
174
  complete_model = complete_model.to(device)
src/priors/normalizing_flow_prior/normalizing_flow_prior.py CHANGED
@@ -48,7 +48,8 @@ class NormalizingFlowPrior(nn.Module):
48
  if weight_path is None:
49
  weight_path = self.param_dict['path_pretrained']
50
  print(' normalizing flow pose prior: loading {}..'.format(weight_path))
51
- pretrained_dict = torch.load(weight_path)['model_state_dict']
 
52
  self.model_inn.load_state_dict(pretrained_dict, strict=True)
53
 
54
  def subnet_fc(self, c_in, c_out):
 
48
  if weight_path is None:
49
  weight_path = self.param_dict['path_pretrained']
50
  print(' normalizing flow pose prior: loading {}..'.format(weight_path))
51
+ device = get_cfg_global_updated().device
52
+ pretrained_dict = torch.load(weight_path, map_location=device)['model_state_dict']
53
  self.model_inn.load_state_dict(pretrained_dict, strict=True)
54
 
55
  def subnet_fc(self, c_in, c_out):