Update app.py
Browse files
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
|
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 |
|