Skip to content
This repository was archived by the owner on Sep 10, 2025. It is now read-only.

Commit ed8ab55

Browse files
authored
Update generate.py
Allow math as fallback
1 parent e3933b2 commit ed8ab55

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

torchchat/generate.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1172,7 +1172,8 @@ def callback(x, *, done_generating=False):
11721172
prof = torch.profiler.profile()
11731173
t0 = time.perf_counter()
11741174
num_tokens_generated = 0
1175-
with torch.nn.attention.sdpa_kernel([self.builder_args.attention_backend]), prof:
1175+
# always allow math as fallback
1176+
with torch.nn.attention.sdpa_kernel([self.builder_args.attention_backend, torch.nn.attention.SDPBackend.MATH]), prof:
11761177
generator_func = self.generate(
11771178
self.model,
11781179
encoded,

0 commit comments

Comments
 (0)