diff --git a/collect_v2.py b/collect_v2.py
index 8ada7aa..5ff44e4 100644
--- a/collect_v2.py
+++ b/collect_v2.py
@@ -83,7 +83,7 @@
                 model="gpt-3.5-turbo",
                 messages=conversation_state,
             )
-            ai_tokens = completion["usage"]["total_tokens"]
+            ai_tokens = ai_completion["usage"]["total_tokens"]
             total_tokens += ai_tokens
             ai_response = ai_completion["choices"][0]["message"]["content"]
             instruct += f"\n[Human] {response}\n[AI] {ai_response}"