From b1bd13616090913bb6bd06b90df9e4ebe83cd76a Mon Sep 17 00:00:00 2001 From: Yong Hoon Shin Date: Wed, 18 Dec 2024 14:31:35 -0800 Subject: [PATCH] Rename _to_sharding_plan -> to_sharding_plan (#2643) Summary: Pull Request resolved: https://github.com/pytorch/torchrec/pull/2643 Renaming the function without underscore, with expectation that other classes may import and use this method too Reviewed By: ge0405 Differential Revision: D67229782 fbshipit-source-id: edb69140c1e9a945d71c908faab7fb7b38f41e88 --- torchrec/distributed/planner/planners.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/torchrec/distributed/planner/planners.py b/torchrec/distributed/planner/planners.py index 1543970d5..101240d60 100644 --- 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)