Use store based barrier only for certain store types. (#49694)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/49694
The store based barrier introduced in
https://github.com/pytorch/pytorch/pull/49419 broke for certain store types.
This is a quick fix to resolve the issues for other store types.
ghstack-source-id: 119006874
Test Plan: 1) waitforbuildbot
Reviewed By: ppwwyyxx, rohan-varma
Differential Revision: D25668404
fbshipit-source-id: 751fb8b229ad6f50ee9c50f63a70de5a91c9eda5
