2024-12-19 nightly release (b1bd136)

pytorch · Dec 19, 2024 · 89b0d91 · 89b0d91
1 parent deb7122
commit 89b0d91
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torchrec/distributed/planner/planners.py b/torchrec/distributed/planner/planners.py
@@ -67,7 +67,7 @@
 from torchrec.distributed.utils import none_throws
 
 
-def _to_sharding_plan(
+def to_sharding_plan(
     sharding_options: List[ShardingOption],
     topology: Topology,
 ) -> ShardingPlan:
@@ -388,7 +388,7 @@ def plan(
                 best_plan = callback(best_plan)
 
             self._best_plan = best_plan
-            sharding_plan = _to_sharding_plan(best_plan, self._topology)
+            sharding_plan = to_sharding_plan(best_plan, self._topology)
 
             end_time = perf_counter()
             for stats in self._stats:
@@ -737,7 +737,7 @@ def plan(
                     best_plan = callback(best_plan)
 
                 self._best_plan = best_plan
-                sharding_plan = _to_sharding_plan(
+                sharding_plan = to_sharding_plan(
                     best_plan, self._topology_groups[group]
                 )
                 best_plans.append(sharding_plan)