zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-08-02 16:56:46 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix data size bug

Browse Source
This commit is contained in:
Jagrit Digani 2024-11-21 12:54:12 -08:00
parent 4640f865cc
commit ed4fb26cb9
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/metal/scaled_dot_product_attention.cpp
View File

@ -330,6 +330,7 @@ void ScaledDotProductAttention::eval_gpu(
size_t str_oH = o.shape(3);
size_t str_oL = o.shape(1) * str_oH;
size_t str_oB = o.shape(2) * str_oL;
size_t data_size = o.shape(0) * str_oB;
array::Flags flags{
/* bool contiguous = */ 1,
@ -339,7 +340,7 @@ void ScaledDotProductAttention::eval_gpu(
o.set_data(
allocator::malloc_or_wait(o.nbytes()),
o.data_size(),
data_size,
{str_oB, str_oH, str_oL, str_oD},
flags);