Suburst commited on
Commit
a690efe
·
verified ·
1 Parent(s): 8f9dc95

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -33,9 +33,9 @@ def load_model(selected_model, model):
33
  global model_path
34
  model_path = os.path.join(model_folder_path, selected_model)
35
  logger.info(f'选择模型路径为{model_path}')
36
- state_dict = torch.load(model_path)
37
  logger.info(f"加载的state_dict大小: {sum([param.numel() for param in state_dict.values()])}")
38
- model.load_state_dict(state_dict, map_location = torch.device('cpu'))
39
  model.eval()
40
  model.to(DEVICE)
41
 
 
33
  global model_path
34
  model_path = os.path.join(model_folder_path, selected_model)
35
  logger.info(f'选择模型路径为{model_path}')
36
+ state_dict = torch.load(model_path,map_location = torch.device('cpu'))
37
  logger.info(f"加载的state_dict大小: {sum([param.numel() for param in state_dict.values()])}")
38
+ model.load_state_dict(state_dict)
39
  model.eval()
40
  model.to(DEVICE)
41