diff --git a/litellm/__init__.py b/litellm/__init__.py
index 591d5873cf..a3727eb2e8 100644
--- a/litellm/__init__.py
+++ b/litellm/__init__.py
@@ -862,6 +862,10 @@ from .llms.vertex_ai_and_google_ai_studio.vertex_ai_anthropic import (
 from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.llama3.transformation import (
     VertexAILlama3Config,
 )
+from .llms.vertex_ai_and_google_ai_studio.vertex_ai_partner_models.ai21.transformation import (
+    VertexAIAi21Config,
+)
+
 from .llms.sagemaker.sagemaker import SagemakerConfig
 from .llms.ollama import OllamaConfig
 from .llms.ollama_chat import OllamaChatConfig