Fix `test_model_parallelism` (#25359)

Commit

2 years ago

Fix `test_model_parallelism` (#25359) * fix * fix --------- Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>