Fix invalid read in masked softmax (#82272) (#82272) (#82405)
Summary:
PEr title, unfortunately testing invalid reads with caching allocator is hard.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82272
Approved by: https://github.com/cpuhrsch
Test Plan:
contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/24d702d38ea4f5b4f8aa9f41e0ae7e589f17b423
Original Phabricator Test Plan:
Imported from GitHub, without a `Test Plan:` line.
Reviewed By: ajtulloch, osalpekar, cpuhrsch
Differential Revision: D38183160
Pulled By: ngimel
fbshipit-source-id: 0ea59868d4829bc540c1277a93daa029519d05b4
Co-authored-by: Natalia Gimelshein (Meta Employee) <ngimel@fb.com>
