Typo: Fix generate signature type hint for attention_backend (#1471)

`attention_backend` is a SDPBackend, not a string
pytorch · Jan 22, 2025 · 025d412 · 025d412
1 parent b2d8f2a
commit 025d412
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/torchchat/generate.py b/torchchat/generate.py
@@ -685,7 +685,7 @@ def generate(
         sequential_prefill=True,
         callback=lambda x: x,
         max_seq_length: int,
-        attention_backend: str = "math",
+        attention_backend: SDPBackend = torch.nn.attention.SDPBackend.MATH,
         seed: Optional[int] = None,
         **sampling_kwargs,
     ) -> torch.Tensor: