diff --git a/examples/models/image_classification/PyPandaDenseNet.py b/examples/models/image_classification/PyPandaDenseNet.py
index 2660dc0b..7a29208a 100644
--- a/examples/models/image_classification/PyPandaDenseNet.py
+++ b/examples/models/image_classification/PyPandaDenseNet.py
@@ -76,9 +76,11 @@ def _create_model(self, scratch: bool, num_classes: int):
     @staticmethod
     def get_knob_config():
         return {
+            'model_class':CategoricalKnob(['densenet']),
             # Learning parameters
             'lr': FixedKnob(0.0001),  ### learning_rate
             'weight_decay': FixedKnob(0.0),
+            'momentum':FixedKnob(0),
             'drop_rate': FixedKnob(0.0),
             'max_epochs': FixedKnob(10),
             'batch_size': CategoricalKnob([32]),
diff --git a/examples/models/image_classification/PyPandaInception.py b/examples/models/image_classification/PyPandaInception.py
index f5982464..12f565aa 100644
--- a/examples/models/image_classification/PyPandaInception.py
+++ b/examples/models/image_classification/PyPandaInception.py
@@ -341,6 +341,7 @@ def get_knob_config():
             # Learning parameters
             'lr':FixedKnob(0.0001), ### learning_rate
             'weight_decay':FixedKnob(0.0),
+            'momentum':FixedKnob(0),
             'drop_rate':FixedKnob(0.0),
             'max_epochs': FixedKnob(30), 
             'batch_size': CategoricalKnob([200]),
diff --git a/examples/models/image_classification/PyPandaResNet.py b/examples/models/image_classification/PyPandaResNet.py
index 2c9ebe49..070f2b50 100644
--- a/examples/models/image_classification/PyPandaResNet.py
+++ b/examples/models/image_classification/PyPandaResNet.py
@@ -64,10 +64,11 @@ def _create_model(self, scratch: bool, num_classes: int):
     @staticmethod
     def get_knob_config():
         return {
-        	'model_class':CategoricalKnob(['resnent101_mnist']),
+        	'model_class':CategoricalKnob(['resnet101_mnist']),
             # Learning parameters
             'lr':FixedKnob(0.0001), ### learning_rate
             'weight_decay':FixedKnob(0.0),
+            'momentum':FixedKnob(0),
             'drop_rate':FixedKnob(0.0),
             'max_epochs': FixedKnob(30), 
             'batch_size': CategoricalKnob([200]),
diff --git a/examples/models/image_classification/PyPandaVgg.py b/examples/models/image_classification/PyPandaVgg.py
index fb0c5765..1e239163 100644
--- a/examples/models/image_classification/PyPandaVgg.py
+++ b/examples/models/image_classification/PyPandaVgg.py
@@ -62,9 +62,11 @@ def _create_model(self, scratch: bool, num_classes: int):
     @staticmethod
     def get_knob_config():
         return {
+            'model_class':CategoricalKnob(['vgg']),
             # Learning parameters
             'lr': FixedKnob(0.0001),  ### learning_rate
             'weight_decay': FixedKnob(0.0),
+            'momentum':FixedKnob(0),
             'drop_rate': FixedKnob(0.0),
             'max_epochs': FixedKnob(1),
             'batch_size': CategoricalKnob([256]),
diff --git a/examples/models/image_classification/PyPandaVgg_SelectiveNet.py b/examples/models/image_classification/PyPandaVgg_SelectiveNet.py
index f55d8dbc..3888b832 100644
--- a/examples/models/image_classification/PyPandaVgg_SelectiveNet.py
+++ b/examples/models/image_classification/PyPandaVgg_SelectiveNet.py
@@ -216,7 +216,7 @@ def train(self,
                 self._gm_optimizer.gm_register(
                     name,
                     f.data.cpu().numpy(),
-                    model_name="PyVGG",
+                    model_name=self._knobs.get("model_class"),
                     hyperpara_list=[
                         self._knobs.get("gm_prior_regularization_a"),
                         self._knobs.get("gm_prior_regularization_b"),
@@ -235,7 +235,7 @@ def train(self,
             self._spl = SPL()
 
         train_dataset = TorchImageDataset(sa_dataset=dataset,
-                                          image_scale_size=128,
+                                          image_scale_size=self._image_size,
                                           norm_mean=self._normalize_mean,
                                           norm_std=self._normalize_std,
                                           is_train=True)
@@ -266,7 +266,14 @@ def train(self,
             optimizer = optim.RMSprop(
                 filter(lambda p: p.requires_grad, self._model.parameters()),
                 lr=self._knobs.get("lr"),
-                weight_decay=self._knobs.get("weight_decay"))
+                weight_decay=self._knobs.get("weight_decay"),
+                momentum=self._knobs.get("momentum"))
+        elif self._knobs.get("optimizer") == "sgd":
+            optimizer = optim.SGD(
+                filter(lambda p: p.requires_grad, self._model.parameters()),
+                lr=self._knobs.get("lr"),
+                weight_decay=self._knobs.get("weight_decay"),
+                momentum=self._knobs.get("momentum"))
         else:
             raise NotImplementedError()
 
@@ -382,7 +389,7 @@ def evaluate(self, dataset_path):
             lazy_load=True)
 
         torch_dataset = TorchImageDataset(sa_dataset=dataset,
-                                          image_scale_size=128,
+                                          image_scale_size=self._image_size,
                                           norm_mean=self._normalize_mean,
                                           norm_std=self._normalize_std,
                                           is_train=False)
@@ -528,9 +535,11 @@ def _create_model(self, scratch: bool, num_classes: int):
     @staticmethod
     def get_knob_config():
         return {
+            'model_class':CategoricalKnob(['vgg_selectivenet']),
             # Learning parameters
             'lr': FixedKnob(0.0001),
             'weight_decay': FixedKnob(0.0),
+            'momentum':FixedKnob(0),
             'drop_rate': FixedKnob(0.0),
             'max_epochs': FixedKnob(10),  # original 5
             'batch_size': CategoricalKnob([96]),  # original 32
diff --git a/singa_easy/models/TorchModel.py b/singa_easy/models/TorchModel.py
index d8499e9f..f3aabc28 100644
--- a/singa_easy/models/TorchModel.py
+++ b/singa_easy/models/TorchModel.py
@@ -249,7 +249,7 @@ def train(self,
                 self._gm_optimizer.gm_register(
                     name,
                     f.data.cpu().numpy(),
-                    model_name="PyVGG",
+                    model_name=self._knobs.get("model_class"),
                     hyperpara_list=[
                         self._knobs.get("gm_prior_regularization_a"),
                         self._knobs.get("gm_prior_regularization_b"),
@@ -292,12 +292,14 @@ def train(self,
             optimizer = optim.RMSprop(
                 filter(lambda p: p.requires_grad, self._model.parameters()),
                 lr=self._knobs.get("lr"),
-                weight_decay=self._knobs.get("weight_decay"))
+                weight_decay=self._knobs.get("weight_decay"),
+                momentum=self._knobs.get("momentum"))
         elif self._knobs.get("optimizer") == "sgd":
             optimizer = optim.SGD(
                 filter(lambda p: p.requires_grad, self._model.parameters()),
                 lr=self._knobs.get("lr"),
-                weight_decay=self._knobs.get("weight_decay"))
+                weight_decay=self._knobs.get("weight_decay"),
+                momentum=self._knobs.get("momentum"))
         else:
             raise NotImplementedError()
 
@@ -548,9 +550,19 @@ def local_explain(self, org_imgs: Image,
                 traceback.print_exc(file=sys.stdout)
 
         if enable_gradcam:
+            if 'densenet' in self._knobs.get("model_class"):
+                model_arch = 'densenet'
+            elif 'alexnet' in self._knobs.get("model_class"):
+                model_arch = 'alexnet'
+            elif 'resnet' in self._knobs.get("model_class"):
+                model_arch = 'resnet'
+            elif 'vgg' in self._knobs.get("model_class"):
+                model_arch = 'vgg'
+            else:
+                raise NameError()
             try:
                 gc = GradCam(model=self._model,
-                             model_arch='vgg',
+                             model_arch=model_arch,
                              target_layer=None,
                              device=self.device)
                 (images, _, _) = utils.dataset.normalize_images(