From bd373b594aa4ddff2de94ecd35713652ab621f8b Mon Sep 17 00:00:00 2001 From: Yong Hoon Shin Date: Fri, 18 Oct 2024 09:49:11 -0700 Subject: [PATCH] fix precommit lint Summary: Precommit (https://github.com/pytorch/torchrec/actions/runs/11396841323/job/31711354638) is failing due to formatting issue Differential Revision: D64606855 --- torchrec/distributed/batched_embedding_kernel.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/torchrec/distributed/batched_embedding_kernel.py b/torchrec/distributed/batched_embedding_kernel.py index 133e2da04..3d157f7f1 100644 --- 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