Update app.py

app.py

@@ -35,7 +35,7 @@ def load_model(selected_model, model):
     logger.info(f'选择模型路径为{model_path}')
     state_dict = torch.load(model_path)
     logger.info(f"加载的state_dict大小: {sum([param.numel() for param in state_dict.values()])}")
-    model.load_state_dict(state_dict)
+    model.load_state_dict(state_dict, map_location = torch.device('cpu'))
     model.eval()
     model.to(DEVICE)