mirror of
https://github.com/ml-explore/mlx.git
synced 2025-12-16 01:49:05 +08:00
Fix the check
This commit is contained in:
@@ -338,7 +338,7 @@ void col_reduce(
|
|||||||
|
|
||||||
// Small col reduce with a single or contiguous reduction axis
|
// Small col reduce with a single or contiguous reduction axis
|
||||||
if (args.non_col_reductions == 1 && args.reduction_size <= 32 &&
|
if (args.non_col_reductions == 1 && args.reduction_size <= 32 &&
|
||||||
args.reduction_stride % 4 == 0) {
|
args.reduction_stride % (16 / in.itemsize()) == 0) {
|
||||||
col_reduce_small(
|
col_reduce_small(
|
||||||
encoder, in, out, reduce_type, axes, plan, std::move(args));
|
encoder, in, out, reduce_type, axes, plan, std::move(args));
|
||||||
return;
|
return;
|
||||||
|
|||||||
Reference in New Issue
Block a user