How to make Cross Entropy Loss work with Cutmix & Mixup?

Hi everyone,
I’m trying to follow the steps of the official tutorial on how to implement both cutmix and mixup during training to perform augmentation but when I start training i get this runtime error from the criterion call.

0D or 1D target tensor expected, multi-target not supported

The following is my training code

    cutmix = v2.CutMix(num_classes=NUM_CLASSES)
    mixup = v2.MixUp(num_classes=NUM_CLASSES)
    cutmix_or_mixup = v2.RandomChoice([cutmix, mixup])
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.learning_rate)    

    for epoch in range(args.epochs):
        train_losses = [] 
        train_acc = 0.0
        total=0
        print(f"[Epoch {epoch+1} / {args.epochs}]")
        
        model.train()
        pbar = tqdm(data_loader)
        for i, (x, y) in enumerate(pbar):
            image = x.to(args.device)
            label = y.to(args.device) 
            image, label = cutmix_or_mixup(image, label)
            optimizer.zero_grad()

            output = model(image)
            label = label.squeeze()
            
            loss = criterion(output, label)
            loss.backward()
            optimizer.step()

            train_losses.append(loss.item())
            total += label.size(0)

            train_acc += acc(output, label)

        epoch_train_loss = np.mean(train_losses)
        epoch_train_acc = train_acc/total

        print(f'Epoch {epoch+1}') 
        print(f'train_loss : {epoch_train_loss}')
        print('train_accuracy : {:.3f}'.format(epoch_train_acc*100))

What am I missing? The tutorial says that I can pass the transformed labels as-is to a loss function like cross entropy.
Tutorial: link

Thanks in advance

If you are passing one-hot encoded labels, make sure they are passed as a floating point tensor. This feature was introduced a few releases ago and allows you to pass “soft” labels to nn.CrossEntropyLoss.

I added label = label.to(torch.float) before the criterion call but I keep getting the same error.

The tutorial works fine for me, so could you post a minimal and executable code snippet reproducing the issue, please?

Apologies for posting on this 3 year old post, but I ran into the same issue today.

Turns out the cause for me was because my network was using an incorrect number of outputs than expected. I had initialized my efficientnet_b0 model with the default number of classes which is 1000, that of ImageNet, when my input data had only 11 classes.

That meant that the cross entropy loss function was comparing tensors of prediction = [num_batch, 1000] against labels = [num_batch, 11] which resulted in the error

RuntimeError: 0D or 1D target tensor expected, multi-target not supported

The code was “working” without the cutup/mixup included because it was comparing prediction= [num_batch, 1000] against labels=[num_batch]

The fix was simple, just provide the num_classes argument to the model constructor

model = efficientnet_b0(weights=None, num_classes=11).to(device)