diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py
index 5291f26916..daa8756da8 100644
--- a/litellm/tests/test_amazing_vertex_completion.py
+++ b/litellm/tests/test_amazing_vertex_completion.py
@@ -814,14 +814,15 @@ def test_vertexai_embedding_embedding_latest():
     try:
         load_vertex_ai_credentials()
         litellm.set_verbose = True
+
         response = embedding(
             model="vertex_ai/text-embedding-004",
-            input=[
-                {"text": "good morning from litellm", "task_type": "RETRIEVAL_DOCUMENT"}
-            ],
+            input=["hi"],
             dimensions=1,
             auto_truncate=True,
+            task_type="RETRIEVAL_QUERY",
         )
+
         assert len(response.data[0]["embedding"]) == 1
         assert response.usage.prompt_tokens > 0
         print(f"response:", response)