fix ao import (#1043)

Summary:

Fix dep issue with torchao

Reviewed By: vkuzo

Differential Revision: D92006625

Author: galrotem

torchtnt/framework/_loop_utils.py

@@ -11,15 +11,28 @@
 
 import torch
 import torch.nn as nn
-import torchao
 from torch.nn.parallel.distributed import DistributedDataParallel
 
 _EXPORT_UTILS_AVAIL = True
 try:
+    from torchao.quantization.pt2e import (
+        move_exported_model_to_eval,
+        move_exported_model_to_train,
+    )
     from torchao.quantization.pt2e.export_utils import model_is_exported
 except Exception:
     _EXPORT_UTILS_AVAIL = False
 
+    def model_is_exported(model: nn.Module) -> bool:
+        return False
+
+    def move_exported_model_to_eval(model: nn.Module) -> None:
+        pass
+
+    def move_exported_model_to_train(model: nn.Module) -> None:
+        pass
+
+
 from torchtnt.utils.progress import Progress
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -102,9 +115,7 @@ def _set_module_training_mode(
             else module
         ):
             move_fn = (
-                torchao.quantization.pt2e.move_exported_model_to_train
-                if mode
-                else torchao.quantization.pt2e.move_exported_model_to_eval
+                move_exported_model_to_train if mode else move_exported_model_to_eval
             )
             # pyre-fixme[6]: For 1st argument expected `GraphModule` but got
             #  `Union[Module, Tensor]`.
@@ -137,9 +148,9 @@ def _reset_module_training_mode(
                 else module
             ):
                 move_fn = (
-                    torchao.quantization.pt2e.move_exported_model_to_train
+                    move_exported_model_to_train
                     if prior_modes[name]
-                    else torchao.quantization.pt2e.move_exported_model_to_eval
+                    else move_exported_model_to_eval
                 )
                 # pyre-fixme[6]: For 1st argument expected `GraphModule` but got
                 #  `Union[Module, Tensor]`.