make model_is_exported oss compatible (#807)

Summary:
Pull Request resolved: #807

`model_is_exported` is only available in nightlies, so let's handle this one as well

bypass-github-export-checks

Reviewed By: diego-urgell

Differential Revision: D56539207

fbshipit-source-id: cc1979bb11d8cec4729cec60ffb68915c378e364

Author: JKSenthil

torchtnt/framework/_loop_utils.py

@@ -13,6 +13,12 @@
 import torch.nn as nn
 from torch.nn.parallel.distributed import DistributedDataParallel
 
+_EXPORT_UTILS_AVAIL = True
+try:
+    from torch.ao.quantization.pt2e.export_utils import model_is_exported
+except Exception:
+    _EXPORT_UTILS_AVAIL = False
+
 from torchtnt.utils.progress import Progress
 
 _logger: logging.Logger = logging.getLogger(__name__)
@@ -66,7 +72,7 @@ def _set_module_training_mode(
         prior_module_train_states[name] = module.training
         is_ddp = isinstance(module, DistributedDataParallel)
 
-        if torch.ao.quantization.pt2e.export_utils.model_is_exported(
+        if _EXPORT_UTILS_AVAIL and model_is_exported(
             module.module if is_ddp else module
         ):
             if mode: