Minor code rewriting to enable torch 1.13

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

Author: 2 people

TrainingExtensions/torch/src/python/aimet_torch/_base/nn/modules/custom.py

@@ -155,9 +155,17 @@ def create_wrapper_module(
 Normalize = create_wrapper_module("Normalize", torch.nn.functional.normalize)
 Pad = create_wrapper_module("Pad", torch.nn.functional.pad)
 GridSample = create_wrapper_module("GridSample", torch.nn.functional.grid_sample)
-ScaledDotProductAttention = create_wrapper_module(
-    "ScaledDotProductAttention", torch.nn.functional.scaled_dot_product_attention
-)
+
+if hasattr(torch.nn.functional, "scaled_dot_product_attention"):
+    ScaledDotProductAttention = create_wrapper_module(
+        "ScaledDotProductAttention", torch.nn.functional.scaled_dot_product_attention
+    )
+elif hasattr(torch.nn.functional, "_scaled_dot_product_attention"):
+    ScaledDotProductAttention = create_wrapper_module(
+        "ScaledDotProductAttention", torch.nn.functional._scaled_dot_product_attention
+    )
+else:
+    ScaledDotProductAttention = None
 
 
 # following modules are for overloaded operators like + and *,

TrainingExtensions/torch/src/python/aimet_torch/v2/experimental/__init__.py

@@ -36,6 +36,10 @@
 # =============================================================================
 
 # pylint: disable=missing-docstring
-from . import export
+try:
+    from . import export
+except ImportError:
+    pass
+
 from . import onnx
 from .quantsim_utils import *

TrainingExtensions/torch/src/python/aimet_torch/v2/quantization/affine/backends/torch_builtins.py

@@ -40,7 +40,11 @@
 from packaging import version
 from typing import Callable, Optional, List, Tuple
 import torch
-import torch.ao.quantization.fx._decomposed
+
+try:
+    import torch.ao.quantization.fx._decomposed
+except ImportError:
+    pass
 from aimet_torch.v2.utils import (
     _is_expandable,
     _ContextManager,