Fix tsan issue (#75528)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75528
ghstack-source-id: 153450975
Test Plan: CI
Reviewed By: fegin
Differential Revision: D35504430
fbshipit-source-id: 0fc299c739d5fb16f4764f32c65d0b904a83013b
(cherry picked from commit 5c7d36b092552b6247758ce09640ccc56e0c47df)