diff --git a/litellm/llms/azure.py b/litellm/llms/azure.py
index 0c8c7f184..6a217bc2c 100644
--- a/litellm/llms/azure.py
+++ b/litellm/llms/azure.py
@@ -715,6 +715,16 @@ class AzureChatCompletion(BaseLLM):
                 model = model
             else:
                 model = None
+
+            ## BASE MODEL CHECK
+            if (
+                model_response is not None
+                and optional_params.get("base_model", None) is not None
+            ):
+                model_response._hidden_params["model"] = optional_params.pop(
+                    "base_model"
+                )
+
             data = {"model": model, "prompt": prompt, **optional_params}
             max_retries = data.pop("max_retries", 2)
             if not isinstance(max_retries, int):
diff --git a/litellm/tests/test_completion_cost.py b/litellm/tests/test_completion_cost.py
index 16ec0602d..f17d5a464 100644
--- a/litellm/tests/test_completion_cost.py
+++ b/litellm/tests/test_completion_cost.py
@@ -297,3 +297,34 @@ def test_whisper_azure():
         5,
     )
     assert cost == expected_cost
+
+
+def test_dalle_3_azure_cost_tracking():
+    litellm.set_verbose = True
+    # model = "azure/dall-e-3-test"
+    # response = litellm.image_generation(
+    #     model=model,
+    #     prompt="A cute baby sea otter",
+    #     api_version="2023-12-01-preview",
+    #     api_base=os.getenv("AZURE_SWEDEN_API_BASE"),
+    #     api_key=os.getenv("AZURE_SWEDEN_API_KEY"),
+    #     base_model="dall-e-3",
+    # )
+    # print(f"response: {response}")
+    response = litellm.ImageResponse(
+        created=1710265780,
+        data=[
+            {
+                "b64_json": None,
+                "revised_prompt": "A close-up image of an adorable baby sea otter. Its fur is thick and fluffy to provide buoyancy and insulation against the cold water. Its eyes are round, curious and full of life. It's lying on its back, floating effortlessly on the calm sea surface under the warm sun. Surrounding the otter are patches of colorful kelp drifting along the gentle waves, giving the scene a touch of vibrancy. The sea otter has its small paws folded on its chest, and it seems to be taking a break from its play.",
+                "url": "https://dalleprodsec.blob.core.windows.net/private/images/3e5d00f3-700e-4b75-869d-2de73c3c975d/generated_00.png?se=2024-03-13T17%3A49%3A51Z&sig=R9RJD5oOSe0Vp9Eg7ze%2FZ8QR7ldRyGH6XhMxiau16Jc%3D&ske=2024-03-19T11%3A08%3A03Z&skoid=e52d5ed7-0657-4f62-bc12-7e5dbb260a96&sks=b&skt=2024-03-12T11%3A08%3A03Z&sktid=33e01921-4d64-4f8c-a055-5bdaffd5e33d&skv=2020-10-02&sp=r&spr=https&sr=b&sv=2020-10-02",
+            }
+        ],
+    )
+    response.usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
+    response._hidden_params = {"model": "dall-e-3", "model_id": None}
+    print(f"response hidden params: {response._hidden_params}")
+    cost = litellm.completion_cost(
+        completion_response=response, call_type="image_generation"
+    )
+    assert cost > 0
diff --git a/litellm/utils.py b/litellm/utils.py
index 3749a6af7..d370b9ce6 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -3841,7 +3841,9 @@ def completion_cost(
                     * n
                 )
             else:
-                raise Exception(f"Model={model} not found in completion cost model map")
+                raise Exception(
+                    f"Model={image_gen_model_name} not found in completion cost model map"
+                )
         # Calculate cost based on prompt_tokens, completion_tokens
         if (
             "togethercomputer" in model