Cross Entropy doesn't work with the specific batch, but works with each sample from this batch
Labels: module: nn, module: cuda, triaged, actionable

### 🐛 Describe the bug

I've found the strange error of cross_entropy loss with one specific batch.

```python

import torch

labels = torch.load('labels.pt').to('cuda:0') # labels.pt is attached
logits = torch.ones(190, 250002, 50).float().to('cuda:0')
loss = torch.nn.functional.cross_entropy(logits, labels)
```

The code above falls down with this CUDA error:

`../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [188,0,0], thread: [32,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [32,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [33,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [34,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [37,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [38,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [40,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [42,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [46,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [183,0,0], thread: [14,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [180,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [180,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [180,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [180,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [18,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [19,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [184,0,0], thread: [20,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [181,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [181,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [181,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [185,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [185,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [185,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [185,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [185,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [12,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [175,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [34,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [37,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [40,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [43,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [45,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [47,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [48,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [173,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [12,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [15,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [16,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [17,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [18,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [21,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [22,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [23,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [24,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [25,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [26,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [28,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [178,0,0], thread: [30,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [174,0,0], thread: [34,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [174,0,0], thread: [38,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [174,0,0], thread: [39,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [174,0,0], thread: [43,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [174,0,0], thread: [48,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [179,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [179,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [179,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [179,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [15,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [17,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [18,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [21,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [22,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [23,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [24,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [25,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [26,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [187,0,0], thread: [27,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [5,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [14,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [15,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [16,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [21,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [23,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [24,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [26,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [27,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [186,0,0], thread: [28,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [177,0,0], thread: [16,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [14,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [17,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [19,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [20,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [21,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [23,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [24,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [27,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [29,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [30,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [172,0,0], thread: [31,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [12,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [14,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [15,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [16,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [182,0,0], thread: [17,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [4,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [7,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [10,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [11,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [12,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [176,0,0], thread: [15,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [1,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [2,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [3,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [6,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [8,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [9,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [12,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [13,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [14,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [19,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [21,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [22,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [23,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [24,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [25,0,0] Assertion `input_index >= 0` failed.
../aten/src/ATen/native/cuda/NLLLoss2d.cu:107: nll_loss2d_forward_kernel: block: [189,0,0], thread: [27,0,0] Assertion `input_index >= 0` failed.`

But, if I pass `cross_entropy` each sample from this batch individually, it works fine. Moreover, the call of `cross_entropy` with entire batch works on cpu.

[labels.pt.zip](https://github.com/pytorch/pytorch/files/12485430/labels.pt.zip)

### Versions

Versions of relevant libraries:
[pip3] imagenetv2-pytorch==0.1
[pip3] lion-pytorch==0.1.2
[pip3] numpy==1.24.4
[pip3] torch==2.0.1
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] Could not collect

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki @ptrblck