diff --git a/tests/test_booster/test_plugin/test_gemini_plugin.py b/tests/test_booster/test_plugin/test_gemini_plugin.py index 57160dfae..a06b2c963 100644 --- a/tests/test_booster/test_plugin/test_gemini_plugin.py +++ b/tests/test_booster/test_plugin/test_gemini_plugin.py @@ -90,7 +90,10 @@ def check_gemini_plugin(init_method: str = 'none', early_stop: bool = True): 'transformers_gpt_double_heads', 'torchaudio_hubert_base', 'torchaudio_wav2vec2_base', 'transformers_t5_for_conditional_generation', 'transformers_t5', 'transformers_t5_encoder_model', 'transformers_vit', 'transformers_vit_for_masked_image_modeling', - 'transformers_vit_for_image_classification' + 'transformers_vit_for_image_classification', 'transformers_chatglm', + 'transformers_chatglm_for_conditional_generation', 'transformers_blip2', + 'transformers_blip2_conditional_gerneration', 'transformers_sam', 'transformers_whisper', + 'transformers_whisperForConditionalGeneration', 'transformers_whisperWhisperForAudioClassification' ]: continue diff --git a/tests/test_lazy/test_models.py b/tests/test_lazy/test_models.py index ecb99e594..18a737fce 100644 --- a/tests/test_lazy/test_models.py +++ b/tests/test_lazy/test_models.py @@ -11,8 +11,9 @@ def test_torchvision_models_lazy_init(subset, default_device): sub_model_zoo = model_zoo.get_sub_registry(subset) for name, entry in sub_model_zoo.items(): # TODO(ver217): lazy init does not support weight norm, skip these models - if name in ('torchaudio_wav2vec2_base', 'torchaudio_hubert_base' - ) or name.startswith('transformers_llama') or name.startswith('transformers_vit'): + if name in ('torchaudio_wav2vec2_base', + 'torchaudio_hubert_base') or name.startswith('transformers_llama') or name.startswith( + ('transformers_vit', 'transformers_blip2')): continue check_lazy_init(entry, verbose=True, default_device=default_device)