Working with Criterion CrossEntropyLoss for multilabel Classification

I know this topic has been covered in the past. But I just wanted to make sure that if it is the correct way of what I’m doing. From my understanding of previous post on this:
→ Label should be torch.long() type
–>We should pass the indicies of the label

    for batch_num, data in enumerate(trainloader, 0):

        # get the images and labels.
        inputs = data["image"].to(device)
        labels = data["label"].to(device)
        labels = labels.long()

        if config.use_amp_grad:

            # forward + backward
            with autocast():
                outputs = net(inputs)
                loss = criterion(outputs, torch.max(labels, 1)[1])

In case you are using nn.CrossEntropyLoss or nn.NLLLoss, then yes: the targets are expected as LongTensors containing the class indices in the range [0, nb_classes-1].

1 Like