From b75d10094d3b279a62e5715cc1623eaa93f347e8 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Mon, 27 Dec 2021 11:31:51 +0800 Subject: [PATCH] fix(mge/elwise): removed back to fp32 mode GitOrigin-RevId: 6d1720f4af7d93ccda2f8c73ebb99a1b25d0150a --- imperative/python/megengine/core/tensor/array_method.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/imperative/python/megengine/core/tensor/array_method.py b/imperative/python/megengine/core/tensor/array_method.py index da6392f32..925934751 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 -- GitLab