Repository: pytorch
Issue: nn.CrossEntropyLoss error out when the sample size is large
Labels: module: nn,module: cuda,module: memory usage,triaged,module: edge cases

### 🐛 Describe the bug

I'm getting the error `RuntimeError: CUDA error: an illegal memory access was encountered` when running `nn.CrossEntropyLoss`. I was able to reproduce the error with the following sample code:
```
import torch
from torch.nn import CrossEntropyLoss
logits_high = 8.8125
logits_low = -6.3125
number_samples = 18423
number_class = 131072
output_class = 50263
# Create logits with uniform distribution between [logits_low, logits_high]
logits = ((logits_high - logits_low) * torch.randn(number_samples, 131072, dtype=torch.bfloat16) + logits_low).to(torch.bfloat16).to(0)
labels = torch.empty(number_samples, dtype=torch.int64)
# Create labels in range [0, output_class)
labels = labels.random_(to=output_class)
labels = labels.to(0)
loss_fct = CrossEntropyLoss()
loss = loss_fct(logits, labels)
print(loss)
```
Running the above script will get the illegal memory error. The tensor size, distribution ranges are picked from real data. The error is sample size related, for example if I change the `number_samples` to 15000 the error will be gone. I'm running on A100 GPU with 40GB memory, this should not be an OOM.

### Versions

```
Collecting environment information...
PyTorch version: 1.12.1+cu113
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.18.2
Libc version: glibc-2.31

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:10)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-1080-aws-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.3.109
CUDA_MODULE_LOADING set to:
GPU models and configuration:
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.2.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.2.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.4
[pip3] sagemaker-pytorch-training==2.7.0
[pip3] torch==1.12.1+cu113
[pip3] torchaudio==0.12.1+cu113
[pip3] torchdata==0.4.1
[pip3] torchnet==0.0.4
[pip3] torchvision==0.13.1+cu113
[conda] magma-cuda113             2.5.2                         1    pytorch
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.23.4           py38h7042d01_1    conda-forge
[conda] sagemaker-pytorch-training 2.7.0                    pypi_0    pypi
[conda] torch                     1.12.1+cu113             pypi_0    pypi
[conda] torchaudio                0.12.1+cu113             pypi_0    pypi
[conda] torchdata                 0.4.1                    pypi_0    pypi
[conda] torchnet                  0.0.4                    pypi_0    pypi
[conda] torchvision               0.13.1+cu113             pypi_0    pypi
```

cc @albanD @mruberry @jbschlosser @walterddr @kshitij12345 @saketh-are @ngimel