Update app.py

app.py

@@ -144,9 +144,9 @@ model_low = ModelCLR(**config["model_config"]).to(device)
 model_median = ModelCLR(**config["model_config"]).to(device)
 model_high = ModelCLR(**config["model_config"]).to(device)
 
-state_dict_low = torch.load(pretrain_model_path_low)
-state_dict_median = torch.load(pretrain_model_path_median)
-state_dict_high = torch.load(pretrain_model_path_high)
+state_dict_low = torch.load(pretrain_model_path_low, map_location=device)
+state_dict_median = torch.load(pretrain_model_path_median, map_location=device)
+state_dict_high = torch.load(pretrain_model_path_high, map_location=device)
 model_low.load_state_dict(state_dict_low)
 model_low.eval()
 model_median.load_state_dict(state_dict_median)