diff --git a/setup.py b/setup.py
index 645207a..b68f365 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
     "install": [
         "setuptools>=38.5.1",
         "librosa>=0.8.0",
-        "onnxruntime>=1.4.0",
+        "onnxruntime",
         "sentencepiece>=0.1.91,!=0.1.92",
         "typeguard>=2.7.0",
         "PyYAML>=5.1.2",
diff --git a/tests/test_optimize.py b/tests/test_optimize.py
index 1cddb4d..75a8f66 100644
--- a/tests/test_optimize.py
+++ b/tests/test_optimize.py
@@ -13,12 +13,12 @@
 
 test_cases = [
     ['encoder', 'transformer', 4, 256, 3, 0, False],
-    ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, False],
-    # ['encoder', 'transformer', 4, 256, 4, 0, True],
-    # ['encoder', 'contextual_block_transformer', 4, 256, 4, 0, True],
-    # ['decoder', 'transformer', 4, 256, 4, 0, True],
-    # ['lm', 'transformer_pe', 4, 256, 4, 0, True],
-    # ['lm', 'transformer', 4, 256, 4, 0, True],
+    # ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, False],
+    ['encoder', 'transformer', 4, 256, 3, 0, True],
+    # ['encoder', 'contextual_block_transformer', 4, 256, 3, 0, True],
+    ['decoder', 'transformer', 4, 256, 0, 6, True],
+    ['lm', 'transformer_pe', 4, 256, 0, 3, True],
+    ['lm', 'transformer', 4, 256, 0, 2, True],
 ]
 
 @pytest.mark.parametrize('model_type, model_name, n_head, h_size, n_att, n_cross_att, use_custom_ort', test_cases)
@@ -26,7 +26,7 @@ def test_optimize(model_type, model_name, n_head, h_size, n_att, n_cross_att, us
     export_dir = model_export.cache_dir / 'test' / \
         model_type / f'cache_{model_name}'
     output_dir = model_export.cache_dir / 'test' / \
-        'optimize' / f'cache_{model_name}'
+        'optimize' / model_type / f'cache_{model_name}'
     
     input_model = glob.glob(os.path.join(export_dir , f'*{model_type}*'))[0]
     model_name = os.path.basename(input_model)