bug fix

app.py

@@ -483,18 +483,29 @@ class ExplainerCheckbox(Component):
         
         data_id = self.gallery.selected_index
         
-        optimized, _, _ = self.experiment.optimize(
-            data_id=data_id.value,
+        opt_output = self.experiment.optimize(
+            data_ids=data_id.value,
             explainer_id=self.default_exp_id,
             metric_id=self.obj_metric,
             direction='maximize',
             sampler=SAMPLE_METHOD,
             n_trials=OPT_N_TRIALS,
         )
+        
 
-        opt_explainer_id = optimized['explainer_id']
-        opt_postprocessor_id = optimized['postprocessor_id']
-
+        def get_str_ppid(pp_obj):
+            return pp_obj.pooling_fn.__class__.__name__ + pp_obj.normalization_fn.__class__.__name__
+        
+        str_id = get_str_ppid(opt_output.postprocessor)
+        for pp_obj, pp_id in zip(*self.experiment.manager.get_postprocessors()):
+            if get_str_ppid(pp_obj) == str_id:
+                opt_postprocessor_id = pp_id
+                break
+
+        opt_explainer_id = max([x['id'] for x in self.groups.info]) + 1
+        opt_output.explainer.model = self.experiment.model
+        self.experiment.manager._explainers.append(opt_output.explainer)
+        self.experiment.manager._explainer_ids.append(opt_explainer_id)
         self.groups.insert_check(self.explainer_name, opt_explainer_id, opt_postprocessor_id)
         self.optimal_exp_id = opt_explainer_id
         checkbox = gr.update(label="Optimized Parameter (Optimal)", interactive=True)
@@ -628,6 +639,8 @@ from torch.utils.data import DataLoader
 from helpers import get_imagenet_dataset, get_torchvision_model, denormalize_image
 
 os.environ['GRADIO_TEMP_DIR'] = '.tmp'
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cpu")
 
 def target_visualizer(x): return dataset.dataset.idx_to_label(x.item())
 
@@ -637,7 +650,7 @@ model, transform = get_torchvision_model('resnet18')
 dataset = get_imagenet_dataset(transform)
 loader = DataLoader(dataset, batch_size=4, shuffle=False)
 experiment1 = AutoExplanationForImageClassification(
-    model=model,
+    model=model.to(device),
     data=loader,
     input_extractor=lambda batch: batch[0],
     label_extractor=lambda batch: batch[-1],
@@ -657,7 +670,7 @@ model, transform = get_torchvision_model('vit_b_16')
 dataset = get_imagenet_dataset(transform)
 loader = DataLoader(dataset, batch_size=4, shuffle=False)
 experiment2 = AutoExplanationForImageClassification(
-    model=model,
+    model=model.to(device),
     data=loader,
     input_extractor=lambda batch: batch[0],
     label_extractor=lambda batch: batch[-1],