fix pre-commit error in sharding (#2755)

Summary:
Pull Request resolved: #2755

# context
* pre-commit test failed: https://github.com/pytorch/torchrec/actions/runs/13416363381/job/37489987778
* after the fix
 {F1975288470}

Reviewed By: iamzainhuda

Differential Revision: D69905081

fbshipit-source-id: 5516a7957b27c581a579c96666c5e96d2522b324

Author: TroyGarden

torchrec/distributed/sharding/cw_sharding.py

@@ -192,7 +192,11 @@ def _shard(
             for i, rank in enumerate(info.param_sharding.ranks):
                 # Remap rank by number of replica groups if 2D parallelism is enabled
                 rank = (
-                    self._env.remap_rank(rank, ShardingType.COLUMN_WISE)  # pyre-ignore[16]
+                    # pyre-ignore[16]
+                    self._env.remap_rank(
+                        rank,
+                        ShardingType.COLUMN_WISE,
+                    )
                     if self._is_2D_parallel
                     else rank
                 )

torchrec/distributed/sharding/grid_sharding.py

@@ -250,7 +250,8 @@ def _shard(
             # pyre-fixme [6]
             for i, rank in enumerate(info.param_sharding.ranks):
                 rank = (
-                    self._env.remap_rank(rank, ShardingType.GRID_SHARD)  # pyre-ignore[16]
+                    # pyre-ignore[16]
+                    self._env.remap_rank(rank, ShardingType.GRID_SHARD)
                     if self._is_2D_parallel
                     else rank
                 )

torchrec/distributed/sharding/tw_sharding.py

@@ -147,7 +147,8 @@ def _shard(
             rank = (
                 # pyre-ignore [16]
                 self._env.remap_rank(
-                    info.param_sharding.ranks[0], ShardingType.TABLE_WISE  # pyre-ignore[16]
+                    info.param_sharding.ranks[0],  # pyre-ignore[16]
+                    ShardingType.TABLE_WISE,
                 )
                 if self._is_2D_parallel
                 else info.param_sharding.ranks[0]