Fix type annotation on logprobs

ProbablyFaiz · Aug 26, 2024 · de5cdff · de5cdff
1 parent b68fa00
commit de5cdff
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/rl/llm/engines.py b/rl/llm/engines.py
@@ -47,7 +47,7 @@ class InferenceOutput(BaseModel):
     prompt: InferenceInput
     text: str
 
-    logprobs: list[dict[str, float]] | None = None
+    logprobs: list[dict[int, float]] | None = None
     metadata: dict[str, Any] = Field(default_factory=dict)