modification

DiseaseSpecific/utils.py

@@ -71,13 +71,13 @@ def load_model(model_path, args, n_ent, n_rel, device):
  model = add_model(args, n_ent, n_rel)
  model.to(device)
  logger.info('Loading saved model from {0}'.format(model_path))
- state = torch.load(model_path)
+ state = torch.load(model_path, map_location=device)
  model_params = state['state_dict']
  params = [(key, value.size(), value.numel()) for key, value in model_params.items()]
  for key, size, count in params:
  logger.info('Key:{0}, Size:{1}, Count:{2}'.format(key, size, count))
  
- model.load_state_dict(model_params, map_location=device)
+ model.load_state_dict(model_params)
  model.eval()
  logger.info(model)