CNN : "Target 5 is out of bounds." error

Using CNN model, I keep getting this “target is out of bounds” error.
What should I fix? my label has 5 classes.

[my model]
스크린샷 2020-05-18 오후 5.14.05

[my code]

epochs = 30
lr = 0.0003
batch_size = 10000

train_idx = np.arange(x_train.size(0))
test_idx = np.arange(x_test.size(0))
optimizer = torch.optim.Adam(model.parameters(),lr)
criterion = nn.CrossEntropyLoss(reduction=‘sum’)

loss_ls = []
for epoch in range(epochs):
model.train()

random.shuffle(train_idx)
x_train = x_train[train_idx]
y_train = y_train[train_idx]
train_loss = 0

for start_idx, end_idx in zip(range(0, x_train.size(0), batch_size),
                              range(batch_size, x_train.size(0)+1, batch_size)):
    
    x_batch = x_train[start_idx : end_idx]
    y_batch = y_train[start_idx : end_idx].long() 
    
    scores = model(x_batch)
    predict = F.softmax(scores, dim=1).argmax(dim = 1)
    
    acc = (predict == y_batch).sum().item() / batch_size
    
    loss = criterion(scores, y_batch)
    train_loss += loss.item()
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    
print('Train epoch : %s,  loss : %s,  accuracy :%.3f'%(epoch+1, train_loss / batch_size, acc))
print('=================================================================================================')

loss_ls.append(train_loss)

if (epoch+1) % 10 == 0:
    model.eval()
    scores = model(x_test)
    predict = F.softmax(scores, dim=1).argmax(dim = 1)
    
    acc = (predict == y_test.long()).sum().item() / y_test.size(0)
    loss = criterion(scores, y_test.long())
    
    print('*************************************************************************************************')
    print('*************************************************************************************************')
    print('Test Epoch : %s, Test Loss : %.03f , Test Accuracy : %.03f'%(epoch+1, loss.item()/y_test.size(0), acc))
    print('*************************************************************************************************')
    print('*************************************************************************************************')

[error msg]

IndexError Traceback (most recent call last)
in ()
30 acc = (predict == y_batch).sum().item() / batch_size
31
—> 32 loss = criterion(scores, y_batch)
33 train_loss += loss.item()
34

3 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
2113 .format(input.size(0), target.size(0)))
2114 if dim == 2:
-> 2115 ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
2116 elif dim == 4:
2117 ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)

IndexError: Target 5 is out of bounds.

I assume you are working on a multi-class classification use case with nn.CrossEntropyLoss as the criterion.
If that’s the case, you would have to make sure that the model output has the shape [batch_size, nb_classes], while the target should have the shape [batch_size] containing the class indices in the range [0, nb_classes-1]. For 5 classes, your target should contain the values [0, 1, 2, 3, 4], while it seems you are passing a 5 to the criterion.