[AOT][CUDAGraphs] torchdynamo -> torch._dynamo (pytorch#87243)

Fixes lingering issues from the torchdynamo -> torch._dynamo migration
Pull Request resolved: pytorch#87243
Approved by: https://github.com/suo, https://github.com/voznesenskym, https://github.com/jansel

torch/cuda/_dynamo_graphs.py

@@ -4,14 +4,14 @@
 from torch.fx.passes.backends.cudagraphs import partition_cudagraphs
 from torch.multiprocessing.reductions import StorageWeakRef
 from torch.utils._pytree import tree_map
-import torchdynamo  # type: ignore[import]
-from torchdynamo.optimizations.training import AOTAutogradStrategy  # type: ignore[import]
+import torch._dynamo  # type: ignore[import]
+from torch._dynamo.optimizations.training import AotAutogradStrategy  # type: ignore[import]
 
 import operator
 from collections import defaultdict
 from typing import Set
 
-# TODO: maybe this should live in torchdynamo instead
+# TODO: maybe this should live in torch._dynamo instead
 
 __all__ = ['aot_autograd_cudagraphs']
 
@@ -140,8 +140,8 @@ def raw_aot_autograd_cudagraphs(model, inputs):
     }
 
     def _wrapped_bw_compiler(*args, **kwargs):
-        # stop TorchDynamo from trying to compile our generated backwards pass
-        return torchdynamo.disable(bw_compiler(*args, **kwargs))  # type: ignore[operator]
+        # stop dynamo from trying to compile our generated backwards pass
+        return torch._dynamo.disable(bw_compiler(*args, **kwargs))  # type: ignore[operator]
 
     bw_compiler = kwargs.get("bw_compiler") or kwargs["fw_compiler"]
     kwargs["bw_compiler"] = _wrapped_bw_compiler
@@ -151,7 +151,7 @@ def _wrapped_bw_compiler(*args, **kwargs):
     return aot_module_simplified(model, **kwargs)
 
 
-class AOTAutogradCudaGraphs(AOTAutogradStrategy):
+class AOTAutogradCudaGraphs(AotAutogradStrategy):
     def candidate(self):
         return raw_aot_autograd_cudagraphs(self.gm, self.example_inputs)