Fixes test_fused_rms_norm.py

fegin · fegin · commit f2a1551663c9 · 2024-09-10T15:43:56.000-07:00
The import is not correct. ghstack-source-id: f06c37f Pull Request resolved: #573
diff --git a/test/test_fused_rms_norm.py b/test/test_fused_rms_norm.py
@@ -11,7 +11,7 @@
     Replicate,
     Shard,
 )
-from torch.distributed._tensor.debug import CommDebugMode
+from torch.distributed.tensor.debug import CommDebugMode
 from torch.testing._internal.common_utils import run_tests
 from torch.testing._internal.distributed._tensor.common_dtensor import (
     DTensorTestBase,

Original file line number	Diff line number	Diff line change
`@@ -11,7 +11,7 @@`
`11`	`11`	`Replicate,`
`12`	`12`	`Shard,`
`13`	`13`	`)`
`14`		`-from torch.distributed._tensor.debug import CommDebugMode`
	`14`	`+from torch.distributed.tensor.debug import CommDebugMode`
`15`	`15`	`from torch.testing._internal.common_utils import run_tests`
`16`	`16`	`from torch.testing._internal.distributed._tensor.common_dtensor import (`
`17`	`17`	`DTensorTestBase,`