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)