diff --git a/vllm/v1/worker/tpu_model_runner.py b/vllm/v1/worker/tpu_model_runner.py index 773c426474fc9..e6c5a8996de6a 100644 --- a/vllm/v1/worker/tpu_model_runner.py +++ b/vllm/v1/worker/tpu_model_runner.py @@ -972,8 +972,7 @@ class TPUModelRunner: def reset_dynamo_cache(self): if self.is_multimodal_model: - assert hasattr(self.model, "language_model") - compiled_model = self.model.language_model.model + compiled_model = self.model.get_language_model().model else: compiled_model = self.model.model if isinstance(compiled_model, TorchCompileWrapperWithCustomDispatcher):