fix: return the out tensor rather then the functions return value (#2361

)
huggingface · Aug 6, 2024 · 29b8d19 · 29b8d19
1 parent dd47a3d
commit 29b8d19
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/server/text_generation_server/layers/attention/cuda.py b/server/text_generation_server/layers/attention/cuda.py
@@ -292,8 +292,7 @@ def attention(
                 )
 
         out = torch.empty_like(q)
-
-        return flash_attn_cuda.fwd(
+        flash_attn_cuda.fwd(
             q,
             k,
             v,
@@ -309,4 +308,5 @@ def attention(
             False,
             0,
             None,
-        )[0]
+        )
+        return out