diff --git a/mmseg/models/segmentors/__init__.py b/mmseg/models/segmentors/__init__.py
index 387c858..f7ede7b 100644
--- a/mmseg/models/segmentors/__init__.py
+++ b/mmseg/models/segmentors/__init__.py
@@ -1,6 +1,11 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from .base import BaseSegmentor
+from .base import BaseSegmentor, ONNXRuntimeSegmentorKN
 from .cascade_encoder_decoder import CascadeEncoderDecoder
 from .encoder_decoder import EncoderDecoder
 
-__all__ = ['BaseSegmentor', 'EncoderDecoder', 'CascadeEncoderDecoder']
+__all__ = [
+    'BaseSegmentor',
+    'ONNXRuntimeSegmentorKN',
+    'EncoderDecoder',
+    'CascadeEncoderDecoder'
+]
diff --git a/tools/deploy_test_kneron.py b/tools/deploy_test_kneron.py
index d1aacdf..aa12f66 100644
--- a/tools/deploy_test_kneron.py
+++ b/tools/deploy_test_kneron.py
@@ -66,6 +66,12 @@ def parse_args() -> argparse.Namespace:
         type=float,
         default=0.5,
         help='Opacity of painted segmentation map. In (0, 1] range.')
+    parser.add_argument(
+        '--shape',
+        type=int,
+        nargs='+',
+        default=None,
+        help='input image height and width.')
     parser.add_argument('--local_rank', type=int, default=0)
     args = parser.parse_args()
     if 'LOCAL_RANK' not in os.environ:
@@ -104,6 +110,28 @@ def main():
         cfg.merge_from_dict(args.cfg_options)
     cfg.model.pretrained = None
     cfg.data.test.test_mode = True
+    if args.shape is not None:
+
+        if len(args.shape) == 1:
+            shape = (args.shape[0], args.shape[0])
+        elif len(args.shape) == 2:
+            shape = (args.shape[1], args.shape[0])
+        else:
+            raise ValueError('invalid input shape')
+
+        test_mode = cfg.model.test_cfg.mode
+        if test_mode == 'slide':
+            warnings.warn(
+                "We suggest you NOT assigning shape when exporting "
+                "slide-mode models. Assigning shape to slide-mode models "
+                "may result in unexpected results. To see which mode the "
+                "model is using, check cfg.model.test_cfg.mode, which "
+                "should be either 'whole' or 'slide'."
+            )
+            cfg.model.test_cfg['crop_size'] = shape
+        else:
+            cfg.test_pipeline[1]['img_scale'] = shape
+        cfg.data.test['pipeline'][1]['img_scale'] = shape
 
     # init distributed env first, since logger depends on the dist info.
     distributed = False
diff --git a/tools/pytorch2onnx_kneron.py b/tools/pytorch2onnx_kneron.py
index 430bdc1..e32b9a5 100644
--- a/tools/pytorch2onnx_kneron.py
+++ b/tools/pytorch2onnx_kneron.py
@@ -2,6 +2,7 @@
 # Original: tools/pytorch2onnx.py, modified by Kneron
 import argparse
 
+import warnings
 import os
 import onnx
 import mmcv
@@ -293,15 +294,24 @@ if __name__ == '__main__':
         else:
             img_scale = cfg.test_pipeline[1]['img_scale']
             input_shape = (1, 3, img_scale[1], img_scale[0])
-    elif len(args.shape) == 1:
-        input_shape = (1, 3, args.shape[0], args.shape[0])
-    elif len(args.shape) == 2:
-        input_shape = (
-            1,
-            3,
-        ) + tuple(args.shape)
     else:
-        raise ValueError('invalid input shape')
+        if test_mode == 'slide':
+            warnings.warn(
+                "We suggest you NOT assigning shape when exporting "
+                "slide-mode models. Assigning shape to slide-mode models "
+                "may result in unexpected results. To see which mode the "
+                "model is using, check cfg.model.test_cfg.mode, which "
+                "should be either 'whole' or 'slide'."
+            )
+        if len(args.shape) == 1:
+            input_shape = (1, 3, args.shape[0], args.shape[0])
+        elif len(args.shape) == 2:
+            input_shape = (
+                1,
+                3,
+            ) + tuple(args.shape)
+        else:
+            raise ValueError('invalid input shape')
 
     # build the model and load checkpoint
     cfg.model.train_cfg = None