Fixed issue in op name retrival (#3383)

Signed-off-by: Chetan Gulecha <[email protected]>
quic · Oct 10, 2024 · 0a0ead8 · 0a0ead8
1 parent 07f2ccf
commit 0a0ead8
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/amp/quantizer_groups.py b/TrainingExtensions/torch/src/python/aimet_torch/amp/quantizer_groups.py
@@ -450,7 +450,7 @@ def find_supported_candidates(quantizer_groups: List[QuantizerGroup],
                                    quantizer)
 
             # pylint: disable=protected-access
-            module = module_name_to_module_dict[quantizer]._module_to_wrap.__class__
+            module = module_name_to_module_dict[quantizer]._module_to_wrap
 
             try:
                 backend_type = aimet_op_to_backend_op_name_map[module.__class__]

diff --git a/TrainingExtensions/torch/src/python/aimet_torch/quantsim_config/quantsim_config.py b/TrainingExtensions/torch/src/python/aimet_torch/quantsim_config/quantsim_config.py
@@ -942,11 +942,11 @@ def generate(self, module: torch.nn.Module, op_type: str) -> Tuple[dict, bool]:
         :return: supported_kernels and per_channel_quantization fields
         """
         supported_kernels = []
-        if module.__class__ in aimet_op_to_backend_op_name_map:
+        if (module.__class__ in aimet_op_to_backend_op_name_map) or (module.__class__.__name__ in aimet_op_to_backend_op_name_map):
             try:
                 backend_type = aimet_op_to_backend_op_name_map[module.__class__]
             except KeyError:
-                backend_type = aimet_op_to_backend_op_name_map.get(module.__class__.__name__)
+                backend_type = aimet_op_to_backend_op_name_map[module.__class__.__name__]
             supported_kernels = self.op_type_supported_kernels.get(backend_type)
 
         if not supported_kernels: