Implement QuantizedTanh (#2975)

Signed-off-by: Daemyung Jang <quic_daemyung@quicinc.com>
quic · May 23, 2024 · 49e8bde · 49e8bde
1 parent 245d12b
commit 49e8bde
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v2/nn/true_quant.py b/TrainingExtensions/torch/src/python/aimet_torch/v2/nn/true_quant.py
@@ -388,6 +388,14 @@ def get_functional_args(self, x):
         return (x, ), {}
 
 
+@QuantizationMixin.implements(nn.Tanh)
+class QuantizedTanh(_QuantizedUnaryOpMixin, nn.Tanh):
+    """ Quantized Tanh """
+
+    def get_functional_args(self, x):
+        return (x,), {}
+
+
 @QuantizationMixin.implements(aimet_ops.Add)
 class QuantizedAdd(_QuantizedBinaryOpMixin, aimet_ops.Add):
     """ Quantized Add """