Reduce memory usage requirement of `test_warp_softmax_64bit_indexing` in `test_nn.py` (#85037)
For reference: #84944
CC @xwang233 @ngimel @ptrblck
Pull Request resolved: https://github.com/pytorch/pytorch/pull/85037
Approved by: https://github.com/ngimel, https://github.com/pmeier
