fix precommit lint

Summary: Precommit (https://github.com/pytorch/torchrec/actions/runs/11396841323/job/31711354638) is failing due to formatting issue Differential Revision: D64606855
pytorch · Oct 18, 2024 · bd373b5 · bd373b5
1 parent d1a2990
commit bd373b5
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py
@@ -476,7 +476,9 @@ def get_sharded_optim_state(
                     momentum_local_shards: List[Shard] = []
                     optimizer_sharded_tensor_metadata: ShardedTensorMetadata
 
-                    optim_state = shard_params.optimizer_states[0][momentum_idx - 1]  # pyre-ignore[16]
+                    optim_state = shard_params.optimizer_states[0][
+                        momentum_idx - 1
+                    ]  # pyre-ignore[16]
                     if (
                         optim_state.nelement() == 1 and state_key != "momentum1"
                     ):  # special handling for backward compatibility, momentum1 is rowwise state for rowwise_adagrad