diff --git a/python/tvm/topi/arm_cpu/injective.py b/python/tvm/topi/arm_cpu/injective.py index 55f47c5dee4d..330144b33fb6 100644 --- a/python/tvm/topi/arm_cpu/injective.py +++ b/python/tvm/topi/arm_cpu/injective.py @@ -17,6 +17,7 @@ # pylint: disable=invalid-name, unused-variable """Schedule for pooling operators""" import tvm +import numpy as np from tvm import te from ..utils import is_empty_shape @@ -67,7 +68,7 @@ def schedule_injective(outs): if list(s[x].op.axis): # do not vectorize for broadcast - (io, ii) = s[x].split(list(s[x].op.axis)[-1], 4) + (io, ii) = s[x].split(list(s[x].op.axis)[-1], 16 // np.dtype(x.dtype).itemsize) s[x].vectorize(ii) tvm.te.schedule.AutoInlineInjective(s)