diff --git a/mlx/backend/metal/kernels/conv.metal b/mlx/backend/metal/kernels/conv.metal index 50a14ad3..24522bb7 100644 --- a/mlx/backend/metal/kernels/conv.metal +++ b/mlx/backend/metal/kernels/conv.metal @@ -30,7 +30,7 @@ template out_pixels *= params->oS[i]; // Set out - out += gid.z * filter_size + gid.y * (params->C); + out += (size_t)gid.z * filter_size + (size_t)gid.y * (params->C); // Coordinates in input int is[N] = {0}; @@ -93,7 +93,8 @@ template out_pixels *= params->oS[i]; // Set out - out += gid.z * filter_size + gid.x * (filter_size / params->C); + out += + (size_t)gid.z * filter_size + (size_t)gid.x * (filter_size / params->C); // Coordinates in input int is[N] = {0};