Update visualization.py

visualization.py

@@ -131,9 +131,9 @@ def genreate_intepriable_output(input,dataset="CUB2011", arch="resnet50",seed=12
     device = torch.device("cpu")
     if folder is None:
         folder = Path(f"tmp/{arch}/{dataset}/{seed}/")
-    model.load_state_dict(torch.load(folder / "Trained_DenseModel.pth"))
-    state_dict = torch.load(folder / f"{model_type}_{n_features}_{n_per_class}_FinetunedModel.pth")
-    selection= torch.load(folder / f"SlDD_Selection_50.pt")
+    model.load_state_dict(torch.load(folder / "Trained_DenseModel.pth", map_location=torch.device('cpu')))
+    state_dict = torch.load(folder / f"{model_type}_{n_features}_{n_per_class}_FinetunedModel.pth", map_location=torch.device('cpu'))
+    selection= torch.load(folder / f"SlDD_Selection_50.pt", map_location=torch.device('cpu'))
     state_dict['linear.selection']=selection
     
     feature_sel, sparse_layer, current_mean, current_std, bias_sparse = extract_sel_mean_std_bias_assignemnt(state_dict)
@@ -370,7 +370,7 @@ def direct_inference(input):
     
     #original
 
-    data_dir=Path.home()/"tmp/Datasets/CUB200/CUB_200_2011/"
+    data_dir=Path("tmp/Datasets/CUB200/CUB_200_2011/")
     classlist=pd.read_csv(data_dir/"classes.txt",sep=' ',names=['cl_id','class_name'])
     output_name=classlist.loc[classlist['cl_id']==output,'class_name'].values[0]
     if concatenated_image[0][1]!=[]: