Fix import failure in PyTorch 1.x (#3468)

Signed-off-by: Kyunggeun Lee <quic_kyunggeu@quicinc.com>
quic · Nov 7, 2024 · c718d18 · c718d18
1 parent 62be1ba
commit c718d18
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/TrainingExtensions/torch/src/python/aimet_torch/utils.py b/TrainingExtensions/torch/src/python/aimet_torch/utils.py
@@ -55,11 +55,15 @@
 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,
 )
+try:
+    from torch.nn.modules.module import _global_backward_pre_hooks
+except ImportError:
+    _global_backward_pre_hooks = None
+
 from torchvision import datasets, transforms
 
 from aimet_common.defs import QuantScheme, QuantizationDataType, MAP_QUANT_SCHEME_TO_PYMO