add batch index in the fwd/bwd#3326

TroyGarden wants to merge 1 commit intometa-pytorch:mainmeta-pytorch/torchrec:mainfrom

TroyGarden:export-D81159416TroyGarden/torchrec:export-D81159416Copy head branch name to clipboard