Minor fix

guestrin-lab · Nov 3, 2024 · 1e1eb4c · 1e1eb4c
1 parent 70f5510
commit 1e1eb4c
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/lotus/models/lm.py b/lotus/models/lm.py
@@ -30,9 +30,7 @@ def __init__(
 
         self.stats: LMStats = LMStats()
 
-    def __call__(
-        self, messages: list[dict[str, str]] | list[list[dict[str, str]]], **kwargs: dict[str, Any]
-    ) -> LMOutput:
+    def __call__(self, messages: list[list[dict[str, str]]], **kwargs: dict[str, Any]) -> LMOutput:
         all_kwargs = {**self.kwargs, **kwargs}
 
         # Set top_logprobs if logprobs requested