diff --git a/imperative/python/megengine/core/tensor/array_method.py b/imperative/python/megengine/core/tensor/array_method.py index da6392f32f020013f37cfaba005d4638a5be16bf..925934751e8a35a0e2782b10ec851a1f9e1e37f6 100644 --- a/imperative/python/megengine/core/tensor/array_method.py +++ b/imperative/python/megengine/core/tensor/array_method.py @@ -50,6 +50,7 @@ def _elwise(*args, mode): if ( mode in ( + _ElwMod.TRUE_DIV, _ElwMod.EXP, _ElwMod.POW, _ElwMod.LOG, @@ -66,7 +67,7 @@ def _elwise(*args, mode): amp._enabled or np.all([np.issubdtype(arg.dtype, np.integer) for arg in args]) ) - or mode in (_ElwMod.TRUE_DIV, _ElwMod.TANH,) + or mode in (_ElwMod.TANH,) and np.all([np.issubdtype(arg.dtype, np.integer) for arg in args]) ): # autocast to FP32 to maintain precision