fix: return the out tensor rather then the functions return value

2025-09-12 04:44:52 +00:00 · 2024-08-05 19:06:55 +00:00 · 2024-08-05 19:06:55 +00:00 · 516b43f006
commit 516b43f006
parent dd47a3dac4
1 changed files with 3 additions and 3 deletions
--- a/server/text_generation_server/layers/attention/cuda.py
+++ b/server/text_generation_server/layers/attention/cuda.py
@ -292,8 +292,7 @@ else:
                )

        out = torch.empty_like(q)
-
-        return flash_attn_cuda.fwd(
+        flash_attn_cuda.fwd(
            q,
            k,
            v,
@ -309,4 +308,5 @@ else:
            False,
            0,
            None,
-        )[0]
+        )
+        return out