Problems with number of classes, adjust CrossEntropyLoss?


When running multi-class classification, I have to set the number of classes as n+1, as setting the number of classes as n returns error.

I attempted to modify CrossEntropyLoss as per instructions on this post, but I must have done something wrong.

I also read this post , but I could not find out a solution.

I have 10 variables split as (X_train,X_test) and 1 target variable split as (Y_train, Y_test). My target/label is a 3 class (1,2,3)

Here’s the code:

class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.layer1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.layer2 = nn.Linear(hidden_size, num_classes)
    def forward(self, x):
        out = self.layer1(x)
        out = self.relu(out)
        out = self.layer2(out)
        return out
tr_latent_X = data_utils.TensorDataset(train_x, train_y)
te_latent_X = data_utils.TensorDataset(test_x, test_y)

x_train_label = Y_train
x_test_label = y_test

train_loader_X =,

test_loader_X =,

If I set the number of dimensions as 3(instead of 4) here later I obtain the error:

vae = NeuralNet(10,2,3)

loss_function = nn.CrossEntropyLoss(reduction='mean',ignore_index=-1)

optimizer = torch.optim.SGD(vae.parameters(), lr = 0.01)

def train(epoch):

    train_loss = 0

    for batch_idx, (data,label) in enumerate(train_loader_X):

        out = vae(data)
        loss = loss_function(out, label)

        train_loss += loss.item()

        if batch_idx % 5000 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader_X.dataset),
                100. * batch_idx / len(train_loader_X), loss.item() / len(data)))
    print('====> Epoch: {} Average loss: {:.4f}'.format(epoch, train_loss / len(train_loader_X.dataset)))

for epoch in range(1, 10):

This is the error I obtained (last part of the long sequenece)

4 frames
/usr/local/lib/python3.6/dist-packages/torch/nn/ in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   1836                          .format(input.size(0), target.size(0)))
   1837     if dim == 2:
-> 1838         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   1839     elif dim == 4:
   1840         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)

RuntimeError: Assertion `cur_target >= 0 && cur_target < n_classes' failed.  at /pytorch/aten/src/THNN/generic/ClassNLLCriterion.c:97

If instead of setting 3 classes, I set 4, then it runs well:

vae = NeuralNet(10,2,4)
loss_function = nn.CrossEntropyLoss(reduction='mean',ignore_index=-1)
optimizer = torch.optim.SGD(vae.parameters(), lr = 0.01)

If anyone could shed some light on this, I would be very grateful.


label -> (0, 1, 2)
Specifically, label will be used as index.
0 <= label < num_classes.

Thanks for your reply @Eta_C but I am not sure how to implement your proposed solution. Where should I place this code?

I’d be grateful if you could elaborate on this.

train_y -= 1
test_y -= 1

1 Like

Thanks, it works!

Should I keep loss function as it is?

loss_function = nn.CrossEntropyLoss(reduction='mean',ignore_index=-1)

or I should remove this part?


Actually, it depends on you. :wink:

well, code runs well with and without it.

Train results look almost alike, so I guess it makes no sense to keep it.

Thanks for your time and help.