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

This commit is contained in:
drbh 2024-08-06 07:49:53 -04:00 committed by GitHub
parent dd47a3dac4
commit 29b8d19cdf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 3 deletions

View File

@ -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