diff --git a/tests/transformation/test_topk_insert.py b/tests/transformation/test_topk_insert.py index a18e63384150f140cb63ec7b438283eb4797266c..a9faac4df0caf973d9aae6430e007eac349a7c43 100644 --- a/tests/transformation/test_topk_insert.py +++ b/tests/transformation/test_topk_insert.py @@ -21,7 +21,7 @@ import pytest export_onnx_path = "test_topk_insert.onnx" -@pytest.mark.parametrize("k", [1, 5, 10]) +@pytest.mark.parametrize("k", [1, 2]) def test_topk_insert(k): tfc = get_test_model_trained("TFC", 1, 1) bo.export_finn_onnx(tfc, (1, 1, 28, 28), export_onnx_path)