Skip to content

GitLab

Unverified Commit 29b8d19c authored by drbh's avatar drbh Committed by GitHub
Browse files

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

parent dd47a3da
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 3 deletions
server/text_generation_server/layers/attention/cuda.py
View file @ 29b8d19c
...@@ -292,8 +292,7 @@ else: ...@@ -292,8 +292,7 @@ else:
) )
out = torch.empty_like(q) out = torch.empty_like(q)
flash_attn_cuda.fwd(
return flash_attn_cuda.fwd(
q, q,
k, k,
v, v,
...@@ -309,4 +308,5 @@ else: ...@@ -309,4 +308,5 @@ else:
False, False,
0, 0,
None, None,
)[0] )
return out
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment