[shard] fix some imports in tests (#73309)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73309
This fix some imports in sharded optimizer tests
ghstack-source-id: 149789036
Test Plan: test_sharded_optim.py
Reviewed By: pritamdamania87
Differential Revision: D34427252
fbshipit-source-id: 41d239889aef7d7740b7bf23556f4af2b3eda01a
(cherry picked from commit c9b8ad8a9e2fdbac2d52c0e6c0c1e8cceefcd836)
