fix: avoid invalid in-place division in CUB-based mean
@@ -59,7 +59,10 @@ cdef ndarray _ndarray_mean(ndarray self, axis, dtype, out, keepdims):
result = cub.cub_reduction(self, cub.CUPY_CUB_SUM, axis, dtype, out,
keepdims)
if result is not None:
+ if result.dtype.kind in ['f', 'c']:
result /= (self.size / result.size)
+ else:
+ result = result / (self.size / result.size)
return result
return _mean(self, axis=axis, dtype=dtype, out=out, keepdims=keepdims)
