diff --git a/tests/test_booster/test_plugin/test_gemini_plugin.py b/tests/test_booster/test_plugin/test_gemini_plugin.py index cf054302e..2e9b24fec 100644 --- a/tests/test_booster/test_plugin/test_gemini_plugin.py +++ b/tests/test_booster/test_plugin/test_gemini_plugin.py @@ -67,6 +67,7 @@ def run_fn(init_method, model_fn, data_gen_fn, output_transform_fn, zero_size, t # TODO(ver217): CI does not support lazy now # @parameterize('init_method', ['lazy', 'none', 'colo']) + @parameterize("subset", [COMMON_MODELS] if IS_FAST_TEST else ["torchvision", "transformers", "diffusers"]) @parameterize("init_method", ["none"]) @parameterize("zero_size", [2])