UNET с функцией потери CrossEntropy

Я пытался обучить UNET с размером ввода как [3,128,128] и соответствующей маской [1,128,128], которая содержит классы напрямую (вместо пикселей она будет содержать номера классов - 1,2). Я пытаюсь решить проблему с двумя классами, поэтому моя маска содержит 1,2 в качестве меток. Теперь я отправляю свои изображения в модель, и размер предсказанных масок составляет [2,128,128]. Теперь, чтобы обучить модель, я выбираю 16 в качестве размера партии. Итак, теперь у меня есть ввод как [16,3,128,128], поэтому прогнозируемая размерность равна [16,2,128,128]. Но у меня есть наземные маски как [16,1,128,128]. Теперь, как я могу применить потерю перекрестной энтропии в Pytorch? Я пробовал следующим образом и получаю следующую ошибку. Не могли бы вы помочь? Заранее спасибо.

lr = 0.1  # 0.1
criterion = nn.CrossEntropyLoss() #nn.L1Loss()
optimizer = optim.SGD(model.parameters(), lr=lr, momentum=0.9, nesterov=True, weight_decay=0.0001)
is_train = True
is_pretrain = False
acc_best = 0
total_epoch = 30
if is_train is True:
    # Training
    for epoch in range(total_epoch):
        model.train()
        tims = time.time()
        for i, (images, labels) in enumerate(train_loader):
            images = Variable(images.permute(0,3,1,2).float().cuda())                                                                                                   
            labels = Variable(labels.type(torch.LongTensor).cuda())
            # Forward + Backward + Optimize
            optimizer.zero_grad()
            outputs = model(images)
            outputs = outputs.type(torch.float)
            print('predictedLabelsType:',outputs[0].type())
            print('ActualLabelsType:',labels[0].type())
            print('shape of predicted outputs:',outputs.shape)
            print('shape of groundtruth masks:',labels.shape)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

Мой вывод выглядит следующим образом:

predictedLabelsType: torch.cuda.FloatTensor
ActualLabelsType: torch.cuda.LongTensor
shape of predicted outputs: torch.Size([16, 2, 128, 128])
shape of groundtruth masks: torch.Size([16, 1, 128, 128])
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-8-b692a8d536a9> in <module>()
     52             print('shape of predicted outputs:',outputs.shape)
     53             print('shape of groundtruth masks:',labels.shape)
---> 54             loss = criterion(outputs, labels)
     55             loss.backward()
     56             optimizer.step()

3 frames
/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in nll_loss(input, target, weight, size_average, ignore_index, reduce, reduction)
   2385         ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2386     elif dim == 4:
-> 2387         ret = torch._C._nn.nll_loss2d(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
   2388     else:
   2389         # dim == 3 or dim > 4

RuntimeError: 1only batches of spatial targets supported (3D tensors) but got targets of size: : [16, 1, 128, 128]

Не могли бы вы подсказать, где ошибка и как этот CrossEntropyLoss() будет работать в Pytorch для сегментации изображений? Что мне здесь не хватает? Я попытался изменить целевой размер как [16,128,128], что приводит к другой ошибке. Большое спасибо!