PyTorch nn.CrossEntropyLoss IndexError: Target 2 is out of bounds
I'm creating a simple 2 class sentiment classifier using bert, but i'm getting an error related to output and label size. I cannot figure out what I'm doing wrong. Below are the required code snippets.
My custom dataset class:
class AmazonReviewsDataset(torch.utils.data.Dataset):
  def __init__(self, df):
    self.df = df
    self.maxlen = 256
    self.tokenizer = DistilBertTokenizer.from_pretrained(&quot;distilbert-base-uncased&quot;)

  def __len__(self):
    return len(self.df)

  def __getitem__(self, index):
    review = self.df['reviews'].iloc[index].split()
    review = ' '.join(review)
    sentiment = int(self.df['sentiment'].iloc[index])

    encodings = self.tokenizer.encode_plus(
        review,
        add_special_tokens=True,
        max_length=self.maxlen,
        padding='max_length',
        truncation=True,
        return_attention_mask=True,
        return_tensors='pt'
    )

    return {
        'input_ids': encodings.input_ids.flatten(),
        'attention_mask': encodings.attention_mask.flatten(),
        'labels': torch.tensor(sentiment, dtype=torch.long)
    }

output of dataloader:
for batch in train_loader:
  print(batch['input_ids'].shape)
  print(batch['attention_mask'].shape)
  print(batch['labels'])
  print(batch['labels'].shape)
  break

torch.Size([32, 256])
torch.Size([32, 256])
tensor([2, 2, 2, 2, 1, 2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 2, 1, 1, 2, 2, 1, 2, 1, 2,
        2, 2, 2, 2, 2, 1, 1, 2])
torch.Size([32])

My nn:
criterion = nn.CrossEntropyLoss().to(device)

class SentimentClassifier(nn.Module):
  def __init__(self):
    super(SentimentClassifier, self).__init__()
    self.distilbert = DistilBertModel.from_pretrained(&quot;distilbert-base-uncased&quot;)
    self.drop0 = nn.Dropout(0.25)
    self.linear1 = nn.Linear(3072, 512)
    self.relu1 = nn.ReLU()
    self.drop1 = nn.Dropout(0.25)
    self.linear2 = nn.Linear(512, 2)
    self.relu2 = nn.ReLU()

  def forward(self, input_ids, attention_mask):
    outputs = self.distilbert(input_ids, attention_mask)
    last_hidden_state = outputs[0]
    pooled_output = torch.cat(tuple([last_hidden_state[:, i] for i in [-4, -3, -2, -1]]), dim=-1)
    x = self.drop0(pooled_output)
    x = self.relu1(self.linear1(x))
    x = self.drop1(x)
    x = self.relu2(self.linear2(x))
    return x

Train loop:
for batch in loop:
    optimizer.zero_grad()
    input_ids = batch['input_ids'].to(device)
    attention_mask = batch['attention_mask'].to(device)
    labels = batch['labels'].to(device)
    output = model(input_ids, attention_mask)
    print(output.size(), labels.size())
    loss = criterion(output, labels) # ERROR
    loss.backward()
    optimizer.step()

Error:
torch.Size([32, 2]) torch.Size([32])

---------------------------------------------------------------------------

IndexError                                Traceback (most recent call last)

&lt;ipython-input-19-6268781f396e&gt; in &lt;module&gt;()
     12     print(output.size(), labels.size())
     13     # output_class = torch.argmax(results, dim=1)
---&gt; 14     loss = criterion(output, labels)
     15     train_loss += loss
     16     loss.backward()

2 frames

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1049         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1050                 or _global_forward_hooks or _global_forward_pre_hooks):
-&gt; 1051             return forward_call(*input, **kwargs)
   1052         # Do not call functions when jit is used
   1053         full_backward_hooks, non_full_backward_hooks = [], []

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/loss.py in forward(self, input, target)
   1119     def forward(self, input: Tensor, target: Tensor) -&gt; Tensor:
   1120         return F.cross_entropy(input, target, weight=self.weight,
-&gt; 1121                                ignore_index=self.ignore_index, reduction=self.reduction)
   1122 
   1123 

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction)
   2822     if size_average is not None or reduce is not None:
   2823         reduction = _Reduction.legacy_get_string(size_average, reduce)
-&gt; 2824     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2825 
   2826 

IndexError: Target 2 is out of bounds.

I read a tutorial which says that, do not use softmax when applying a nn.CrossEntropyLoss, as I have 2 classes. What is wrong can anyone guide me! Thank you!

You have two classes, which means the maximum target label is 1 not 2 because the classes are indexed from 0. You essentially have to subtract 1 to your labels tensor, such that class n°1 is assigned the value 0, and class n°2 value 1.
In turn the labels of the batch you printed would look like:
tensor([1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1, 0, 0, 1, 1, 0, 1, 0, 1,
        1, 1, 1, 1, 1, 0, 0, 1])

