From 9f828b519b5781d7d612b43e90c673d94af07652 Mon Sep 17 00:00:00 2001 From: Seokjun An Date: Fri, 3 Jan 2025 00:37:56 +0900 Subject: [PATCH] Fix AMP v2 slowdown issue Signed-off-by: Seokjun An --- .../torch/src/python/aimet_torch/v2/amp/utils.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/TrainingExtensions/torch/src/python/aimet_torch/v2/amp/utils.py b/TrainingExtensions/torch/src/python/aimet_torch/v2/amp/utils.py index c461862b031..68ecd4aa6fe 100644 --- a/TrainingExtensions/torch/src/python/aimet_torch/v2/amp/utils.py +++ b/TrainingExtensions/torch/src/python/aimet_torch/v2/amp/utils.py @@ -108,6 +108,15 @@ def forward(self, x): return super().forward(x) + @contextmanager + def compute_encodings(self): + if not self.enabled: + yield + return + + with super().compute_encodings(): + yield + @classmethod def from_v2_quantizer(cls, qtzr: Optional[QuantizeDequantize]) -> Union['_V1DisabledQuantizer', '_V1QuantizeDequantize']: