diff --git a/mpa/cls/exporter.py b/mpa/cls/exporter.py
index 8b2eeefd..a28d6da0 100644
--- a/mpa/cls/exporter.py
+++ b/mpa/cls/exporter.py
@@ -70,7 +70,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
 
         data = self.get_fake_input(cfg)
         fake_img = data['img'].unsqueeze(0)
-        
+
         precision = kwargs.pop('precision', 'FP32')
         logger.info(f'Model will be exported with precision {precision}')
 
diff --git a/mpa/cls/stage.py b/mpa/cls/stage.py
index 18e7bc6c..bbba29ad 100644
--- a/mpa/cls/stage.py
+++ b/mpa/cls/stage.py
@@ -99,10 +99,6 @@ def is_mmov_model(k, v):
             if cfg.model.get('multilabel', False) or cfg.model.get('hierarchical', False):
                 cfg.model.head.pop('topk', None)
 
-        # Other hyper-parameters
-        if cfg.get('hyperparams', False):
-            self.configure_hyperparams(cfg, training, **kwargs)
-
         return cfg
 
     @staticmethod
@@ -169,7 +165,7 @@ def configure_task(cfg, training, model_meta=None, **kwargs):
 
         model_tasks, dst_classes = None, None
         model_classes, data_classes = [], []
-        train_data_cfg = Stage.get_train_data_cfg(cfg)
+        train_data_cfg = Stage.get_data_cfg(cfg, "train")
         if isinstance(train_data_cfg, list):
             train_data_cfg = train_data_cfg[0]
 
@@ -290,18 +286,6 @@ def configure_task(cfg, training, model_meta=None, **kwargs):
                 cfg.model.head.num_old_classes = len(old_classes)
         return model_tasks, dst_classes
 
-    @staticmethod
-    def configure_hyperparams(cfg, training, **kwargs):
-        hyperparams = kwargs.get('hyperparams', None)
-        if hyperparams is not None:
-            bs = hyperparams.get('bs', None)
-            if bs is not None:
-                cfg.data.samples_per_gpu = bs
-
-            lr = hyperparams.get('lr', None)
-            if lr is not None:
-                cfg.optimizer.lr = lr
-
     def _put_model_on_gpu(self, model, cfg):
         if torch.cuda.is_available():
             model = model.cuda()
diff --git a/mpa/cls/trainer.py b/mpa/cls/trainer.py
index f99abfc8..77841e27 100644
--- a/mpa/cls/trainer.py
+++ b/mpa/cls/trainer.py
@@ -104,10 +104,14 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
 
         # prepare data loaders
         datasets = datasets if isinstance(datasets, (list, tuple)) else [datasets]
-        train_data_cfg = Stage.get_train_data_cfg(cfg)
-        drop_last = train_data_cfg.drop_last if train_data_cfg.get('drop_last', False) else False
-
-        # updated to adapt list of datasets for the 'train'
+        train_data_cfg = Stage.get_data_cfg(cfg, "train")
+        otx_dataset = train_data_cfg.get("otx_dataset", None)
+        drop_last = False
+        dataset_len = len(otx_dataset) if otx_dataset else 0
+        # if task == h-label & dataset size is bigger than batch size
+        if train_data_cfg.get("hierarchical_info", None) and dataset_len > cfg.data.get("samples_per_gpu", 2):
+            drop_last = True
+        # updated to adapt list of dataset for the 'train'
         data_loaders = []
         sub_loaders = []
         for ds in datasets:
@@ -179,7 +183,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
             cfg.optimizer_config.pop('type')
             optimizer_config = opt_hook(
                 **cfg.optimizer_config, **fp16_cfg, distributed=self.distributed)
-        elif self.distributed and 'type' not in cfg.optimizer_config:
+        elif self.distributed and "type" not in cfg.optimizer_config:
             optimizer_config = DistOptimizerHook(**cfg.optimizer_config)
         else:
             optimizer_config = cfg.optimizer_config
@@ -197,7 +201,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
         for hook in cfg.get('custom_hooks', ()):
             runner.register_hook_from_cfg(hook)
 
-        validate = True if cfg.data.get('val', None) else False
+        validate = True if cfg.data.get("val", None) else False
         # register eval hooks
         if validate:
             val_dataset = build_dataset(cfg.data.val, dict(test_mode=True))
@@ -218,26 +222,26 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
             runner.resume(cfg.resume_from)
         elif cfg.get('load_from', False):
             if self.distributed:
-                runner.load_checkpoint(cfg.load_from, map_location=f'cuda:{cfg.gpu_ids[0]}')
+                runner.load_checkpoint(cfg.load_from, map_location=f"cuda:{cfg.gpu_ids[0]}")
             else:
                 runner.load_checkpoint(cfg.load_from)
         runner.run(data_loaders, cfg.workflow)
 
-        logger.info(f'called train_worker() distributed={self.distributed}, validate=True')
+        logger.info(f"called train_worker() distributed={self.distributed}, validate=True")
 
         # Save outputs
-        output_ckpt_path = osp.join(cfg.work_dir, 'best_model.pth'
-                                    if osp.exists(osp.join(cfg.work_dir, 'best_model.pth'))
-                                    else 'latest.pth')
+        output_ckpt_path = osp.join(cfg.work_dir, "best_model.pth"
+                                    if osp.exists(osp.join(cfg.work_dir, "best_model.pth"))
+                                    else "latest.pth")
         return dict(final_ckpt=output_ckpt_path)
 
     def _modify_cfg_for_distributed(self, model, cfg):
         nn.SyncBatchNorm.convert_sync_batchnorm(model)
 
-        if cfg.dist_params.get('linear_scale_lr', False):
+        if cfg.dist_params.get("linear_scale_lr", False):
             new_lr = len(cfg.gpu_ids) * cfg.optimizer.lr
-            logger.info(f'enabled linear scaling rule to the learning rate. \
-                changed LR from {cfg.optimizer.lr} to {new_lr}')
+            logger.info(f"enabled linear scaling rule to the learning rate. \
+                changed LR from {cfg.optimizer.lr} to {new_lr}")
             cfg.optimizer.lr = new_lr
 
     @staticmethod
diff --git a/mpa/det/exporter.py b/mpa/det/exporter.py
index f3570bfb..dcf47d60 100644
--- a/mpa/det/exporter.py
+++ b/mpa/det/exporter.py
@@ -34,6 +34,7 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
 
         output_path = os.path.join(cfg.work_dir, 'export')
         os.makedirs(output_path, exist_ok=True)
+
         model = build_detector(cfg.model)
         if model_ckpt:
             load_checkpoint(model=model, filename=model_ckpt, map_location='cpu')
diff --git a/mpa/det/incremental/stage.py b/mpa/det/incremental/stage.py
index 341817fe..e7cfbcff 100644
--- a/mpa/det/incremental/stage.py
+++ b/mpa/det/incremental/stage.py
@@ -95,7 +95,7 @@ def configure_classes(self, cfg, task_adapt_type, task_adapt_op):
 
     def configure_task_data_pipeline(self, cfg, model_classes, data_classes):
         # Trying to alter class indices of training data according to model class order
-        tr_data_cfg = self.get_train_data_cfg(cfg)
+        tr_data_cfg = self.get_data_cfg(cfg, "train")
         class_adapt_cfg = dict(type='AdaptClassLabels', src_classes=data_classes, dst_classes=model_classes)
         pipeline_cfg = tr_data_cfg.pipeline
         for i, op in enumerate(pipeline_cfg):
@@ -121,7 +121,7 @@ def configure_task_cls_incr(self, cfg, task_adapt_type, org_model_classes, model
             self.configure_ema(cfg)
             self.configure_val_interval(cfg)
         else:
-            src_data_cfg = self.get_train_data_cfg(cfg)
+            src_data_cfg = self.get_data_cfg(cfg, "train")
             src_data_cfg.pop('old_new_indices', None)
 
     def configure_bbox_head(self, cfg, model_classes):
@@ -136,7 +136,7 @@ def configure_bbox_head(self, cfg, model_classes):
         # TODO Remove this part
         # This is not related with patching bbox head
         # This might be useless when semisl using MPADetDataset
-        tr_data_cfg = self.get_train_data_cfg(cfg)
+        tr_data_cfg = self.get_data_cfg(cfg, "train")
         if tr_data_cfg.type != 'MPADetDataset':
             tr_data_cfg.img_ids_dict = self.get_img_ids_for_incr(cfg, org_model_classes, model_classes)
             tr_data_cfg.org_type = tr_data_cfg.type
@@ -228,7 +228,7 @@ def get_img_ids_for_incr(cfg, org_model_classes, model_classes):
         new_classes = np.setdiff1d(model_classes, org_model_classes).tolist()
         old_classes = np.intersect1d(org_model_classes, model_classes).tolist()
 
-        src_data_cfg = self.get_train_data_cfg(cfg)
+        src_data_cfg = self.get_data_cfg(cfg, "train")
 
         ids_old, ids_new = [], []
         data_cfg = cfg.data.test.copy()
diff --git a/mpa/det/inferrer.py b/mpa/det/inferrer.py
index 0d042959..66623506 100644
--- a/mpa/det/inferrer.py
+++ b/mpa/det/inferrer.py
@@ -34,9 +34,9 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
         """
         self._init_logger()
         mode = kwargs.get('mode', 'train')
-        eval = kwargs.get('eval', False)
-        dump_features = kwargs.get('dump_features', False)
-        dump_saliency_map = kwargs.get('dump_saliency_map', False)
+        eval = kwargs.pop("eval", False)
+        dump_features = kwargs.pop("dump_features", False)
+        dump_saliency_map = kwargs.pop("dump_saliency_map", False)
         if mode not in self.mode:
             return {}
 
@@ -83,7 +83,7 @@ def infer(self, cfg, eval=False, dump_features=False, dump_saliency_map=False):
             input_source = cfg.get('input_source')
             logger.info(f'Inferring on input source: data.{input_source}')
             if input_source == 'train':
-                src_data_cfg = self.get_train_data_cfg(cfg)
+                src_data_cfg = self.get_data_cfg(cfg, input_source)
             else:
                 src_data_cfg = cfg.data[input_source]
             data_cfg.test_mode = src_data_cfg.get('test_mode', False)
diff --git a/mpa/det/semisl/inferrer.py b/mpa/det/semisl/inferrer.py
index 89c2c46e..133db256 100644
--- a/mpa/det/semisl/inferrer.py
+++ b/mpa/det/semisl/inferrer.py
@@ -82,7 +82,7 @@ def infer(self, cfg, eval=False, dump_features=False, dump_saliency_map=False):
             input_source = cfg.get('input_source')
             logger.info(f'Inferring on input source: data.{input_source}')
             if input_source == 'train':
-                src_data_cfg = self.get_train_data_cfg(cfg)
+                src_data_cfg = self.get_data_cfg(cfg, "train")
             else:
                 src_data_cfg = cfg.data[input_source]
             data_cfg.test_mode = src_data_cfg.get('test_mode', False)
diff --git a/mpa/det/semisl/stage.py b/mpa/det/semisl/stage.py
index f92003fa..2d8a8d9b 100644
--- a/mpa/det/semisl/stage.py
+++ b/mpa/det/semisl/stage.py
@@ -47,7 +47,7 @@ def configure_task_cls_incr(self, cfg, task_adapt_type, org_model_classes, model
             self.configure_task_adapt_hook(cfg, org_model_classes, model_classes)
             self.configure_val_interval(cfg)
         else:
-            src_data_cfg = self.get_train_data_cfg(cfg)
+            src_data_cfg = self.get_data_cfg(cfg, "train")
             src_data_cfg.pop('old_new_indices', None)
 
     @staticmethod
diff --git a/mpa/det/semisl/trainer.py b/mpa/det/semisl/trainer.py
index cf0316f1..ae4286f0 100644
--- a/mpa/det/semisl/trainer.py
+++ b/mpa/det/semisl/trainer.py
@@ -153,10 +153,10 @@ def train_worker(gpu, target_classes, datasets, cfg, distributed=False,
         model = build_detector(cfg.model)
         model.CLASSES = target_classes
         # Do clustering for SSD model
-        # TODO[JAEGUK]: Temporal Disable cluster_anchors for SSD model
+        # TODO[JAEGUK]: Temporary disable cluster_anchors for SSD model
         # if hasattr(cfg.model, 'bbox_head') and hasattr(cfg.model.bbox_head, 'anchor_generator'):
         #     if getattr(cfg.model.bbox_head.anchor_generator, 'reclustering_anchors', False):
-        #         train_cfg = Stage.get_train_data_cfg(cfg)
+        #         train_cfg = Stage.get_data_cfg(cfg, "train")
         #         train_dataset = train_cfg.get('otx_dataset', None)
         #         cfg, model = cluster_anchors(cfg, train_dataset, model)
         train_detector(
diff --git a/mpa/det/stage.py b/mpa/det/stage.py
index da842ae5..999a6f1c 100644
--- a/mpa/det/stage.py
+++ b/mpa/det/stage.py
@@ -119,7 +119,7 @@ def configure_data(self, cfg, data_cfg, training, **kwargs):
             cfg.data.train.type = super_type
         if training:
             if 'dataset' in cfg.data.train:
-                train_cfg = self.get_train_data_cfg(cfg)
+                train_cfg = self.get_data_cfg(cfg, "train")
                 train_cfg.otx_dataset = cfg.data.train.pop('otx_dataset', None)
                 train_cfg.labels = cfg.data.train.get('labels', None)
                 train_cfg.data_classes = cfg.data.train.pop('data_classes', None)
diff --git a/mpa/det/trainer.py b/mpa/det/trainer.py
index daf50f3b..f5dfbb55 100644
--- a/mpa/det/trainer.py
+++ b/mpa/det/trainer.py
@@ -65,6 +65,13 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
 
         # Data
         datasets = [build_dataset(cfg.data.train)]
+        cfg.data.val.samples_per_gpu = cfg.data.get('samples_per_gpu', 1)
+
+        # FIXME: scale_factors is fixed at 1 even batch_size > 1 in simple_test_mask
+        # Need to investigate, possibly due to OpenVINO
+        if 'roi_head' in model_cfg.model:
+            if 'mask_head' in model_cfg.model.roi_head:
+                cfg.data.val.samples_per_gpu = 1
 
         if hasattr(cfg, 'hparams'):
             if cfg.hparams.get('adaptive_anchor', False):
@@ -113,10 +120,10 @@ def run(self, model_cfg, model_ckpt, data_cfg, **kwargs):
             self._modify_cfg_for_distributed(model, cfg)
 
         # Do clustering for SSD model
-        # TODO[JAEGUK]: Temporal Disable cluster_anchors for SSD model
+        # TODO[JAEGUK]: Temporary disable cluster_anchors for SSD model
         # if hasattr(cfg.model, 'bbox_head') and hasattr(cfg.model.bbox_head, 'anchor_generator'):
         #     if getattr(cfg.model.bbox_head.anchor_generator, 'reclustering_anchors', False):
-        #         train_cfg = Stage.get_train_data_cfg(cfg)
+        #         train_cfg = Stage.get_data_cfg(cfg, "train")
         #         train_dataset = train_cfg.get('otx_dataset', None)
         #         cfg, model = cluster_anchors(cfg, train_dataset, model)
 
diff --git a/mpa/modules/hooks/checkpoint_hook.py b/mpa/modules/hooks/checkpoint_hook.py
index 43579edc..79645729 100644
--- a/mpa/modules/hooks/checkpoint_hook.py
+++ b/mpa/modules/hooks/checkpoint_hook.py
@@ -51,17 +51,17 @@ def __init__(self,
     def after_train_epoch(self, runner):
         if not self.by_epoch or not self.every_n_epochs(runner, self.interval):
             return
-        if hasattr(runner, 'save_ckpt'):
-            if runner.save_ckpt:
-                if runner.save_ema_model:
-                    backup_model = runner.model
-                    runner.model = runner.ema_model
-                runner.logger.info(f'Saving checkpoint at {runner.epoch + 1} epochs')
-                if self.sync_buffer:
-                    allreduce_params(runner.model.buffers())
-                self._save_checkpoint(runner)
-                if runner.save_ema_model:
-                    runner.model = backup_model
+        if hasattr(runner, 'save_ckpt') and runner.save_ckpt:
+            if hasattr(runner, 'save_ema_model') and runner.save_ema_model:
+                backup_model = runner.model
+                runner.model = runner.ema_model
+            runner.logger.info(f'Saving checkpoint at {runner.epoch + 1} epochs')
+            if self.sync_buffer:
+                allreduce_params(runner.model.buffers())
+            self._save_checkpoint(runner)
+            if hasattr(runner, 'save_ema_model') and runner.save_ema_model:
+                runner.model = backup_model
+                runner.save_ema_model = False
             runner.save_ckpt = False
 
     @master_only
diff --git a/mpa/modules/hooks/recording_forward_hooks.py b/mpa/modules/hooks/recording_forward_hooks.py
index c3e3df6a..3d45725d 100644
--- a/mpa/modules/hooks/recording_forward_hooks.py
+++ b/mpa/modules/hooks/recording_forward_hooks.py
@@ -34,8 +34,9 @@ class BaseRecordingForwardHook(ABC):
             print(hook.records)
     Args:
         module (torch.nn.Module): The PyTorch module to be registered in forward pass
+        fpn_idx (int, optional): The layer index to be processed if the model is a FPN.
+                                  Defaults to 0 which uses the largest feature map from FPN.
     """
-
     def __init__(self, module: torch.nn.Module, fpn_idx: int = 0) -> None:
         self._module = module
         self._handle = None
@@ -259,6 +260,12 @@ def func(self, feature_map: Union[torch.Tensor, Sequence[torch.Tensor]], fpn_idx
         """
         Generate the class-wise saliency maps using Recipro-CAM and then normalizing to (0, 255).
 
+        Args:
+            feature_map (Union[torch.Tensor, List[torch.Tensor]]): feature maps from backbone or list of feature maps
+                                                                    from FPN.
+            fpn_idx (int, optional): The layer index to be processed if the model is a FPN.
+                                      Defaults to 0 which uses the largest feature map from FPN.
+
         Returns:
             torch.Tensor: Class-wise Saliency Maps. One saliency map per each class - [batch, class_id, H, W]
         """
diff --git a/mpa/modules/models/classifiers/sam_classifier.py b/mpa/modules/models/classifiers/sam_classifier.py
index ba744228..a445d8e7 100644
--- a/mpa/modules/models/classifiers/sam_classifier.py
+++ b/mpa/modules/models/classifiers/sam_classifier.py
@@ -106,9 +106,9 @@ def state_dict_hook(module, state_dict, *args, **kwargs):
         if backbone_type == 'OTXMobileNetV3':
             for k, v in state_dict.items():
                 if k.startswith('backbone'):
-                    k = k.replace('backbone.', '')
+                    k = k.replace('backbone.', '', 1)
                 elif k.startswith('head'):
-                    k = k.replace('head.', '')
+                    k = k.replace('head.', '', 1)
                     if '3' in k:  # MPA uses "classifier.3", OTX uses "classifier.4". Convert for OTX compatibility.
                         k = k.replace('3', '4')
                         if module.multilabel and not module.is_export:
@@ -118,9 +118,9 @@ def state_dict_hook(module, state_dict, *args, **kwargs):
         elif backbone_type == 'OTXEfficientNet':
             for k, v in state_dict.items():
                 if k.startswith('backbone'):
-                    k = k.replace('backbone.', '')
+                    k = k.replace('backbone.', '', 1)
                 elif k.startswith('head'):
-                    k = k.replace('head', 'output')
+                    k = k.replace('head', 'output', 1)
                     if not module.hierarchical and not module.is_export:
                         k = k.replace('fc', 'asl')
                         v = v.t()
@@ -129,7 +129,7 @@ def state_dict_hook(module, state_dict, *args, **kwargs):
         elif backbone_type == 'OTXEfficientNetV2':
             for k, v in state_dict.items():
                 if k.startswith('backbone'):
-                    k = k.replace('backbone.', '')
+                    k = k.replace('backbone.', '', 1)
                 elif k == 'head.fc.weight':
                     k = k.replace('head.fc', 'model.classifier')
                     if not module.hierarchical and not module.is_export:
diff --git a/mpa/modules/models/heads/custom_cls_head.py b/mpa/modules/models/heads/custom_cls_head.py
index c805d4ce..e6e52324 100644
--- a/mpa/modules/models/heads/custom_cls_head.py
+++ b/mpa/modules/models/heads/custom_cls_head.py
@@ -2,6 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 #
 
+import torch
+import torch.nn.functional as F
+
 from mmcls.models.builder import HEADS
 from mmcls.models.heads import LinearClsHead
 from .non_linear_cls_head import NonLinearClsHead
@@ -83,6 +86,17 @@ def loss(self, cls_score, gt_label, feature=None):
         losses['loss'] = loss
         return losses
 
+    def simple_test(self, img):
+        """Test without augmentation."""
+        cls_score = self.fc(img)
+        if isinstance(cls_score, list):
+            cls_score = sum(cls_score) / float(len(cls_score))
+        if torch.onnx.is_in_onnx_export():
+            return cls_score
+        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
+
+        return self.post_process(pred)
+
     def forward_train(self, x, gt_label):
         cls_score = self.fc(x)
         losses = self.loss(cls_score, gt_label, feature=x)
diff --git a/mpa/modules/models/heads/custom_hierarchical_linear_cls_head.py b/mpa/modules/models/heads/custom_hierarchical_linear_cls_head.py
index ef123367..9b43738b 100644
--- a/mpa/modules/models/heads/custom_hierarchical_linear_cls_head.py
+++ b/mpa/modules/models/heads/custom_hierarchical_linear_cls_head.py
@@ -36,6 +36,10 @@ def __init__(self,
         self.hierarchical_info = kwargs.pop('hierarchical_info', None)
         assert self.hierarchical_info
         super(CustomHierarchicalLinearClsHead, self).__init__(loss=loss)
+        if self.hierarchical_info['num_multiclass_heads'] + \
+            self.hierarchical_info['num_multilabel_classes'] == 0:
+            raise ValueError(
+                'Invalid classification heads configuration')
         self.compute_multilabel_loss = False
         if self.hierarchical_info['num_multilabel_classes'] > 0:
             self.compute_multilabel_loss = build_loss(multilabel_loss)
@@ -111,14 +115,21 @@ def simple_test(self, img):
         for i in range(self.hierarchical_info['num_multiclass_heads']):
             multiclass_logit = cls_score[:, self.hierarchical_info['head_idx_to_logits_range'][i][0]:
                                             self.hierarchical_info['head_idx_to_logits_range'][i][1]]
+            if not torch.onnx.is_in_onnx_export():
+                multiclass_logit = torch.softmax(multiclass_logit, dim=1)
             multiclass_logits.append(multiclass_logit)
-        multiclass_logits = torch.cat(multiclass_logits, dim=1)
-        multiclass_pred = torch.softmax(multiclass_logits, dim=1) if multiclass_logits is not None else None
+        multiclass_pred = torch.cat(multiclass_logits, dim=1) if multiclass_logits else None
 
         if self.compute_multilabel_loss:
             multilabel_logits = cls_score[:, self.hierarchical_info['num_single_label_classes']:]
-            multilabel_pred = torch.sigmoid(multilabel_logits) if multilabel_logits is not None else None
-            pred = torch.cat([multiclass_pred, multilabel_pred], axis=1)
+            if not torch.onnx.is_in_onnx_export():
+                multilabel_pred = torch.sigmoid(multilabel_logits) if multilabel_logits is not None else None
+            else:
+                multilabel_pred = multilabel_logits
+            if multiclass_pred is not None:
+                pred = torch.cat([multiclass_pred, multilabel_pred], axis=1)
+            else:
+                pred = multilabel_pred
         else:
             pred = multiclass_pred
 
diff --git a/mpa/modules/models/heads/custom_hierarchical_non_linear_cls_head.py b/mpa/modules/models/heads/custom_hierarchical_non_linear_cls_head.py
index 12f44f4f..1d52c9c4 100644
--- a/mpa/modules/models/heads/custom_hierarchical_non_linear_cls_head.py
+++ b/mpa/modules/models/heads/custom_hierarchical_non_linear_cls_head.py
@@ -42,6 +42,10 @@ def __init__(self,
         self.hierarchical_info = kwargs.pop('hierarchical_info', None)
         assert self.hierarchical_info
         super(CustomHierarchicalNonLinearClsHead, self).__init__(loss=loss)
+        if self.hierarchical_info['num_multiclass_heads'] + \
+            self.hierarchical_info['num_multilabel_classes'] == 0:
+            raise ValueError(
+                'Invalid classification heads configuration')
         self.compute_multilabel_loss = False
         if self.hierarchical_info['num_multilabel_classes'] > 0:
             self.compute_multilabel_loss = build_loss(multilabel_loss)
@@ -138,14 +142,21 @@ def simple_test(self, img):
         for i in range(self.hierarchical_info['num_multiclass_heads']):
             multiclass_logit = cls_score[:, self.hierarchical_info['head_idx_to_logits_range'][i][0]:
                                             self.hierarchical_info['head_idx_to_logits_range'][i][1]]
+            if not torch.onnx.is_in_onnx_export():
+                multiclass_logit = torch.softmax(multiclass_logit, dim=1)
             multiclass_logits.append(multiclass_logit)
-        multiclass_logits = torch.cat(multiclass_logits, dim=1)
-        multiclass_pred = torch.softmax(multiclass_logits, dim=1) if multiclass_logits is not None else None
+        multiclass_pred = torch.cat(multiclass_logits, dim=1) if multiclass_logits else None
 
         if self.compute_multilabel_loss:
             multilabel_logits = cls_score[:, self.hierarchical_info['num_single_label_classes']:]
-            multilabel_pred = torch.sigmoid(multilabel_logits) if multilabel_logits is not None else None
-            pred = torch.cat([multiclass_pred, multilabel_pred], axis=1)
+            if not torch.onnx.is_in_onnx_export():
+                multilabel_pred = torch.sigmoid(multilabel_logits) if multilabel_logits is not None else None
+            else:
+                multilabel_pred = multilabel_logits
+            if multiclass_pred is not None:
+                pred = torch.cat([multiclass_pred, multilabel_pred], axis=1)
+            else:
+                pred = multilabel_pred
         else:
             pred = multiclass_pred
 
diff --git a/mpa/modules/models/heads/custom_multi_label_linear_cls_head.py b/mpa/modules/models/heads/custom_multi_label_linear_cls_head.py
index b29648a6..530ae1b6 100644
--- a/mpa/modules/models/heads/custom_multi_label_linear_cls_head.py
+++ b/mpa/modules/models/heads/custom_multi_label_linear_cls_head.py
@@ -81,9 +81,9 @@ def simple_test(self, img):
         cls_score = self.fc(img) * self.scale
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = torch.sigmoid(cls_score) if cls_score is not None else None
         if torch.onnx.is_in_onnx_export():
-            return pred
+            return cls_score
+        pred = torch.sigmoid(cls_score) if cls_score is not None else None
         pred = list(pred.detach().cpu().numpy())
         return pred
 
diff --git a/mpa/modules/models/heads/custom_multi_label_non_linear_cls_head.py b/mpa/modules/models/heads/custom_multi_label_non_linear_cls_head.py
index 755635da..3253f9f0 100644
--- a/mpa/modules/models/heads/custom_multi_label_non_linear_cls_head.py
+++ b/mpa/modules/models/heads/custom_multi_label_non_linear_cls_head.py
@@ -105,9 +105,9 @@ def simple_test(self, img):
         cls_score = self.classifier(img) * self.scale
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = torch.sigmoid(cls_score) if cls_score is not None else None
         if torch.onnx.is_in_onnx_export():
-            return pred
+            return cls_score
+        pred = torch.sigmoid(cls_score) if cls_score is not None else None
         pred = list(pred.detach().cpu().numpy())
         return pred
 
diff --git a/mpa/modules/models/heads/non_linear_cls_head.py b/mpa/modules/models/heads/non_linear_cls_head.py
index 7ce46d97..047d3d0b 100644
--- a/mpa/modules/models/heads/non_linear_cls_head.py
+++ b/mpa/modules/models/heads/non_linear_cls_head.py
@@ -76,9 +76,9 @@ def simple_test(self, img):
         cls_score = self.classifier(img)
         if isinstance(cls_score, list):
             cls_score = sum(cls_score) / float(len(cls_score))
-        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
         if torch.onnx.is_in_onnx_export():
-            return pred
+            return cls_score
+        pred = F.softmax(cls_score, dim=1) if cls_score is not None else None
         pred = list(pred.detach().cpu().numpy())
         return pred
 
diff --git a/mpa/seg/inferrer.py b/mpa/seg/inferrer.py
index 06d361e1..cb820acc 100644
--- a/mpa/seg/inferrer.py
+++ b/mpa/seg/inferrer.py
@@ -58,7 +58,7 @@ def infer(self, cfg, dump_features=False):
         input_source = cfg.get('input_source', 'test')
         self.logger.info(f'Inferring on input source: data.{input_source}')
         if input_source == 'train':
-            src_data_cfg = Stage.get_train_data_cfg(cfg)
+            src_data_cfg = Stage.get_data_cfg(cfg, "train")
         else:
             src_data_cfg = cfg.data[input_source]
         data_cfg = cfg.data.test.copy()
diff --git a/mpa/seg/semisl/inferrer.py b/mpa/seg/semisl/inferrer.py
index a1d2662a..663ee417 100644
--- a/mpa/seg/semisl/inferrer.py
+++ b/mpa/seg/semisl/inferrer.py
@@ -66,7 +66,7 @@ def infer(self, cfg, dump_features=False):
         input_source = cfg.get('input_source', 'test')
         self.logger.info(f'Inferring on input source: data.{input_source}')
         if input_source == 'train':
-            src_data_cfg = Stage.get_train_data_cfg(cfg)
+            src_data_cfg = Stage.get_data_cfg(cfg, "train")
         else:
             src_data_cfg = cfg.data[input_source]
         data_cfg = cfg.data.test.copy()
diff --git a/mpa/seg/stage.py b/mpa/seg/stage.py
index 395524ac..1f3b171d 100644
--- a/mpa/seg/stage.py
+++ b/mpa/seg/stage.py
@@ -26,7 +26,6 @@ def configure(self, model_cfg, model_ckpt, data_cfg, training=True, **kwargs):
         self.configure_ckpt(cfg, model_ckpt, kwargs.get('pretrained', None))
         self.configure_data(cfg, data_cfg, training)
         self.configure_task(cfg, training, **kwargs)
-        self.configure_hyperparams(cfg, training, **kwargs)
 
         return cfg
 
@@ -81,7 +80,7 @@ def configure_data(self, cfg, data_cfg, training):
             if cfg.data.get('val', False):
                 self.validate = True
         # Dataset
-        src_data_cfg = Stage.get_train_data_cfg(cfg)
+        src_data_cfg = Stage.get_data_cfg(cfg, "train")
         for mode in ['train', 'val', 'test']:
             if src_data_cfg.type == 'MPASegDataset' and cfg.data.get(mode, False):
                 if cfg.data[mode]['type'] != 'MPASegDataset':
@@ -90,18 +89,6 @@ def configure_data(self, cfg, data_cfg, training):
                     cfg.data[mode]['type'] = 'MPASegDataset'
                     cfg.data[mode]['org_type'] = org_type
 
-    def configure_hyperparams(self, cfg, training, **kwargs):
-        if 'hyperparams' in cfg:
-            hyperparams = kwargs.get('hyperparams', None)
-            if hyperparams is not None:
-                bs = hyperparams.get('bs', None)
-                if bs is not None:
-                    cfg.data.samples_per_gpu = bs
-
-                lr = hyperparams.get('lr', None)
-                if lr is not None:
-                    cfg.optimizer.lr = lr
-
     def configure_task(self, cfg, training, **kwargs):
         """Adjust settings for task adaptation
         """
diff --git a/mpa/stage.py b/mpa/stage.py
index 6fbf5e12..ed84fbf7 100644
--- a/mpa/stage.py
+++ b/mpa/stage.py
@@ -268,19 +268,20 @@ def get_model_meta(cfg):
         return meta
 
     @staticmethod
-    def get_train_data_cfg(cfg):
-        if 'dataset' in cfg.data.train:  # Concat|RepeatDataset
-            dataset = cfg.data.train.dataset
+    def get_data_cfg(cfg, subset):
+        assert subset in ["train", "val", "test"], f"Unknown subset:{subset}"
+        if 'dataset' in cfg.data[subset]:  # Concat|RepeatDataset
+            dataset = cfg.data[subset].dataset
             while hasattr(dataset, 'dataset'):
                 dataset = dataset.dataset
             return dataset
         else:
-            return cfg.data.train
+            return cfg.data[subset]
 
     @staticmethod
     def get_data_classes(cfg):
         data_classes = []
-        train_cfg = Stage.get_train_data_cfg(cfg)
+        train_cfg = Stage.get_data_cfg(cfg, "train")
         if 'data_classes' in train_cfg:
             data_classes = list(train_cfg.pop('data_classes', []))
         elif 'classes' in train_cfg:
diff --git a/mpa/utils/convert_keys.py b/mpa/utils/convert_keys.py
index c61708f8..e69de29b 100644
--- a/mpa/utils/convert_keys.py
+++ b/mpa/utils/convert_keys.py
@@ -1,69 +0,0 @@
-# Copyright (C) 2022 Intel Corporation
-# SPDX-License-Identifier: Apache-2.0
-#
-
-import torch
-
-from mpa.utils.logger import get_logger
-
-logger = get_logger()
-
-
-def convert_keys(name, path, new_path=None):
-    if not path:
-        return path
-
-    if not new_path:
-        new_path = path[:-3] + 'converted.pth'
-    if torch.cuda.is_available():
-        ckpt = torch.load(path)
-    else:
-        ckpt = torch.load(path, map_location='cpu')
-
-    new_ckpt = {}
-
-    if 'state_dict' in ckpt.keys():
-        state_dict = ckpt.pop('state_dict')
-        new_ckpt['state_dict'] = {}
-        new_state_dict = new_ckpt['state_dict']
-        for k, v in ckpt.items():
-            new_ckpt[k] = v
-    else:
-        state_dict = ckpt
-        new_state_dict = new_ckpt
-
-    if name == 'OTEMobileNetV3':
-        for k, v in state_dict.items():
-            if k.startswith('classifier.'):
-                new_state_dict['head.'+k] = v
-            elif not k.startswith('backbone.') and not k.startswith('head.'):
-                new_state_dict['backbone.'+k] = v
-            else:
-                new_state_dict[k] = v
-    elif name == 'OTEEfficientNet':
-        for k, v in state_dict.items():
-            # if 'output.' in k:
-            #     new_state_dict['head.'+k[7:]] = v
-            # else:
-            #     new_state_dict['backbone.'+k] = v
-            if k.startswith('output.'):
-                v = v.t()
-                if 'asl' in k:
-                    new_state_dict['head.fc' + k[10:]] = v
-                else:
-                    new_state_dict['head.'+k[7:]] = v
-            elif not k.startswith('backbone.') and not k.startswith('head.'):
-                if 'activ.' in k:
-                    k = k.replace('activ.', 'activate.')
-                new_state_dict['backbone.'+k] = v
-            else:
-                new_state_dict[k] = v
-    else:
-        raise ValueError(f'Not supported model - {name}')
-
-    if new_state_dict == state_dict:
-        logger.info('conversion is not required.')
-        return path
-
-    torch.save(new_ckpt, new_path)
-    return new_path
diff --git a/recipes/stages/_base_/schedules/1cycle.py b/recipes/stages/_base_/schedules/1cycle.py
index bbf00435..15ff1e6f 100644
--- a/recipes/stages/_base_/schedules/1cycle.py
+++ b/recipes/stages/_base_/schedules/1cycle.py
@@ -2,7 +2,7 @@
 
 lr_config = dict(
     policy='OneCycle',
-    pct_start=0.2,
+    pct_start=0.200001,
     div_factor=100,
     final_div_factor=1000
 )
diff --git a/recipes/stages/classification/train.yaml b/recipes/stages/classification/train.yaml
index f697db5a..48880a3d 100644
--- a/recipes/stages/classification/train.yaml
+++ b/recipes/stages/classification/train.yaml
@@ -6,12 +6,12 @@ _base_: [
     '../_base_/schedules/cos_anneal.py',
 ]
 
-optimizer: 
+optimizer:
     lr: 0.03
     momentum: 0.9
 
-runner: 
+runner:
     max_epochs: 10
 
-evaluation: 
+evaluation:
     metric: ['accuracy']