From 49e8bde16c3167981829081eefbb3b58dd531dda Mon Sep 17 00:00:00 2001
From: Daemyung Jang <quic_daemyung@quicinc.com>
Date: Thu, 23 May 2024 15:54:41 +0900
Subject: [PATCH] Implement QuantizedTanh (#2975)

Signed-off-by: Daemyung Jang <quic_daemyung@quicinc.com>
---
 .../torch/src/python/aimet_torch/v2/nn/true_quant.py      | 8 ++++++++
 1 file changed, 8 insertions(+)

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
index 7b4f95c80e..0c95dee4d5 100644
--- 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 """