[ZeroRedundancyOptimizer] bounding the multiple gpus unit test to 4 gpus, hardcoded values (#54788)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/53322, the test has some hardcoded values to check that the sharding works as expected, and was not used beyond 4 gpus prior
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54788
Reviewed By: mrshenli
Differential Revision: D27483078
Pulled By: blefaudeux
fbshipit-source-id: 63fe072c41e1601925af23d8fb1ea3f4729b2044
