Update

[ghstack-poisoned]
pytorch · Jun 18, 2024 · abc6032 · abc6032
2 parents 63d8f9a + 3eba8cd
commit abc6032
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/torch/distributed/_tensor/examples/comm_mode_features_example.py b/torch/distributed/_tensor/examples/comm_mode_features_example.py
@@ -5,6 +5,8 @@
 import torch
 from torch.distributed._tensor import DeviceMesh
 from torch.distributed._tensor.debug import CommDebugMode
+from torch.distributed._tensor.examples.comm_mode_features_example_argparser import args
+
 from torch.distributed.tensor.parallel import (
     ColwiseParallel,
     parallelize_module,