diff --git a/src/llmtuner/model/utils/visual.py b/src/llmtuner/model/utils/visual.py
index b8696096..1f770861 100644
--- a/src/llmtuner/model/utils/visual.py
+++ b/src/llmtuner/model/utils/visual.py
@@ -8,7 +8,7 @@ from ...extras.logging import get_logger
 
 
 if TYPE_CHECKING:
-    from transformers import LlavaConfig, PretrainedConfig, PreTrainedModel, LlavaForConditionalGeneration
+    from transformers import LlavaConfig, PretrainedConfig, PreTrainedModel
 
     from ...hparams import ModelArguments