fix torch package issue caused by heterogenous planner (pytorch#1905)

Summary: Fix torch package issue caused context S410933 Reviewed By: IvanKobzarev Differential Revision: D56360099
gnahzg · Apr 19, 2024 · bd26adc · bd26adc
1 parent d0f7729
commit bd26adc
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 9 deletions.
diff --git a/torchrec/distributed/planner/__init__.py b/torchrec/distributed/planner/__init__.py
@@ -21,9 +21,6 @@
     - automatically building and selecting an optimized sharding plan.
 """
 
-from torchrec.distributed.planner.planners import (  # noqa
-    EmbeddingShardingPlanner,
-    HeteroEmbeddingShardingPlanner,  # noqa
-)
+from torchrec.distributed.planner.planners import EmbeddingShardingPlanner  # noqa
 from torchrec.distributed.planner.types import ParameterConstraints, Topology  # noqa
 from torchrec.distributed.planner.utils import bytes_to_gb, sharder_name  # noqa
diff --git a/torchrec/distributed/shard.py b/torchrec/distributed/shard.py
@@ -15,11 +15,8 @@
 from torch.distributed._composable.contract import contract
 from torchrec.distributed.comm import get_local_size
 from torchrec.distributed.model_parallel import get_default_sharders
-from torchrec.distributed.planner import (
-    EmbeddingShardingPlanner,
-    HeteroEmbeddingShardingPlanner,
-    Topology,
-)
+from torchrec.distributed.planner import EmbeddingShardingPlanner, Topology
+from torchrec.distributed.planner.planners import HeteroEmbeddingShardingPlanner
 from torchrec.distributed.sharding_plan import (
     get_module_to_default_sharders,
     ParameterShardingGenerator,