[Distributed] Fix new token's shape #3028

Job	Run time
compile-gguf (macos-14)	4m 30s
gather-models-gpu	15s
runner-aoti (16-core-ubuntu)	4m 36s
runner-aoti (macos-14-xlarge)	3m 1s
gather-models-cpu	9s
runner-et (16-core-ubuntu)	6m 26s
runner-et (macos-14-xlarge)	4m 53s
test-gguf-util (macos-14)	4m 1s
test-mps (macos-m1-stable)	2m 6s
test-mps-dtype (macos-m1-stable)	1m 19s
test-tinystories-executorch (16-core-ubuntu)	19m 15s
test-torchao-experimental (macos-14-xlarge)	5m 41s
test-tinystories-executorch (macos-14-xlarge)	17m 23s
torchchat-command-load-test (macos-14)	1m 47s
test-build-runner-et-android / linux-job	13m 40s
test-cpu-compile (x86_64, stories15M)	16m 9s
test-cpu-aoti (x86_64, stories15M)	17m 32s
test-cpu-compile (aarch64, stories15M)	9m 24s
test-cpu-aoti (aarch64, stories15M)	9m 41s
test-cpu-eval-sanity-check (x86_64, stories15M)	9m 4s
test-cpu-eval-sanity-check-float16 (x86_64, stories15M)	29m 16s
test-cpu-eval-sanity-check (aarch64, stories15M)	11m 27s
test-cpu-eval-sanity-check-float16 (aarch64, stories15M)	6m 27s
test-cpu-eval-sanity-check-float32 (x86_64, stories15M)	36m 19s
test-cpu-eval-sanity-check-float32 (aarch64, stories15M)	8m 41s
test-gpu-aoti-float16 (cuda, stories15M) / linux-job	12m 3s
test-gpu-aoti-bfloat16 (cuda, stories15M) / linux-job	13m 3s
test-gpu-aoti-float32 (cuda, stories15M) / linux-job	12m 1s
test-gpu-compile (cuda, stories15M) / linux-job	23m 18s
test-gpu-eval-sanity-check (cuda, stories15M) / linux-job	11m 33s
	5h 15m 0s

Provide feedback