Fix mypy

guestrin-lab · Nov 3, 2024 · b05784a · b05784a
1 parent 0c86fc8
commit b05784a
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lotus/models/lm.py b/lotus/models/lm.py
@@ -50,6 +50,9 @@ def __call__(
         return LMOutput(outputs=outputs, logprobs=logprobs)
 
     def _update_stats(self, response: ModelResponse):
+        if not hasattr(response, "usage"):
+            return
+
         self.stats.total_usage.prompt_tokens += response.usage.prompt_tokens
         self.stats.total_usage.completion_tokens += response.usage.completion_tokens
         self.stats.total_usage.total_tokens += response.usage.total_tokens