diff --git a/docs/docs/guides/tracking/costs.md b/docs/docs/guides/tracking/costs.md
index 8bcddeb2e0c..bedca15aa17 100644
--- a/docs/docs/guides/tracking/costs.md
+++ b/docs/docs/guides/tracking/costs.md
@@ -1,9 +1,5 @@
 # Costs
 
-:::info
-Custom costs are accessible via Python and REST queries. UI uptake is under development and expected to be complete by middle of October 2024
-:::
-
 ## Adding a custom cost
 
 You can add a custom cost by using the [`add_cost`](/reference/python-sdk/weave/trace/weave.trace.weave_client#method-add_cost) method.
diff --git a/docs/docs/reference/gen_notebooks/custom_model_cost.md b/docs/docs/reference/gen_notebooks/custom_model_cost.md
index fc7bf819c2a..093286e9f06 100644
--- a/docs/docs/reference/gen_notebooks/custom_model_cost.md
+++ b/docs/docs/reference/gen_notebooks/custom_model_cost.md
@@ -92,6 +92,7 @@ class YourModel(Model):
             "usage": {
                 "input_tokens": prompt_tokens,
                 "output_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
             },
             "model": "your_model_name",
             "output": prediction,
diff --git a/docs/notebooks/custom_model_cost.ipynb b/docs/notebooks/custom_model_cost.ipynb
index 87b0187593d..def49145f9f 100644
--- a/docs/notebooks/custom_model_cost.ipynb
+++ b/docs/notebooks/custom_model_cost.ipynb
@@ -12,7 +12,6 @@
     "---\n",
     "docusaurus_head_meta::end -->\n",
     "\n",
-    
     "<!--- @wandbcode{prompt-optim-notebook} -->"
    ]
   },
@@ -196,6 +195,7 @@
     "            \"usage\": {\n",
     "                \"input_tokens\": prompt_tokens,\n",
     "                \"output_tokens\": completion_tokens,\n",
+    "                \"total_tokens\": prompt_tokens + completion_tokens,\n",
     "            },\n",
     "            \"model\": \"your_model_name\",\n",
     "            \"output\": prediction,\n",