[bugfix][relay] fix wrong calculate logic about celu (apache#14796)

* fix celu * add test case * Update test_forward.py * Update pytorch.py
leemiracle · May 8, 2023 · e01cb47 · e01cb47
1 parent d6e0f1d
commit e01cb47
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py
@@ -968,8 +968,9 @@ def celu(self, inputs, input_types):
         data = inputs[0]
         dtype = input_types[0]
         alpha = _expr.const(float(inputs[1]), dtype=dtype)
-        return alpha * _op.nn.relu(
-            _expr.const(1, dtype=dtype) - _op.exp(data / alpha)
+        zero = _op.const(0, dtype)
+        return alpha * _op.minimum(
+            zero, _op.exp(data / alpha) - _expr.const(1, dtype=dtype)
         ) + _op.nn.relu(data)
 
     def gelu(self, inputs, input_types):

diff --git a/tests/python/frontend/pytorch/test_forward.py b/tests/python/frontend/pytorch/test_forward.py
@@ -786,6 +786,8 @@ def test_forward_celu():
     verify_model(torch.nn.CELU(alpha=0.3).eval(), input_data=input_data)
     verify_model(torch.nn.CELU(alpha=1.0).eval(), input_data=input_data)
     verify_model(torch.nn.CELU(alpha=1.3).eval(), input_data=input_data)
+    input_data = torch.tensor([-1.0, 2.0], dtype=torch.float32)
+    verify_model(torch.nn.CELU().eval(), input_data=input_data)
 
 
 @tvm.testing.uses_gpu