Fix minor bug in has_hooks (#3460)

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>

TrainingExtensions/torch/src/python/aimet_torch/utils.py

@@ -54,6 +54,12 @@
 import torch.nn
 import torch
 from torch.utils.data import DataLoader, Dataset
+from torch.nn.modules.module import (
+    _global_backward_pre_hooks,
+    _global_backward_hooks,
+    _global_forward_pre_hooks,
+    _global_forward_hooks,
+)
 from torchvision import datasets, transforms
 
 from aimet_common.defs import QuantScheme, QuantizationDataType, MAP_QUANT_SCHEME_TO_PYMO
@@ -543,14 +549,15 @@ def get_input_shape_batch_size(data_loader):
 
 def has_hooks(module: torch.nn.Module):
     """ Returns True if the module uses hooks. """
-
-    for hooks in (module._forward_pre_hooks,                       # pylint: disable=protected-access
-                  module._forward_hooks, module._backward_hooks):  # pylint: disable=protected-access
-        if hooks:
-            logger.warning("The specified model has registered hooks which might break winnowing")
-            return True
-    return False
-
+    # pylint: disable=protected-access
+    return module._backward_hooks or\
+           module._backward_pre_hooks or\
+           module._forward_hooks or\
+           module._forward_pre_hooks or\
+           _global_backward_pre_hooks or\
+           _global_backward_hooks or\
+           _global_forward_hooks or\
+           _global_forward_pre_hooks
 
 def get_one_positions_in_binary_mask(mask):
     """

TrainingExtensions/torch/src/python/aimet_torch/winnow/mask_propagation_winnower.py

@@ -73,7 +73,9 @@ def __init__(self, model: torch.nn.Module, input_shape: Tuple,
         """
 
         super().__init__(list_of_modules_to_winnow, reshape, in_place, verbose)
-        model.apply(has_hooks)
+
+        if any(has_hooks(module) for module in model.modules()):
+            logger.warning("The specified model has registered hooks which might break winnowing")
 
         debug_level = logger.getEffectiveLevel()
         logger.debug("Current log level: %s", debug_level)