Форма достоверной информации в мультиклассовой сегментации изображений с помощью pytorch

Я работаю над изображениями ячеек 128 x 128 x 3 и хочу разделить их на 5 классов, включая фоны. Сначала я сделал целевые изображения 128 x 128 и значениями в {0,1,2,3,4}. Но я обнаружил, что мне нужно сделать свою целевую наземную истину как 5-канальное изображение, и все значения равны 0 или 1: если пиксель имеет 1 в n-м канале, то он должен быть отнесен к n-му классу.

Но когда я запустил свою модель в модель Unet, которую я разветвил из GitHub, я обнаружил ошибку при вычислении кросс-энтропийной потери.

Изначально я установил количество каналов на входе равным 3, а количество классов на выходе равным 5. И размер пакета = 2.

Вот мои коды:

for i, (x, y) in batch_iter:
    input, target = x.to(self.device), y.to(self.device)  # send to device (GPU or CPU)
    self.optimizer.zero_grad()  # zerograd the parameters
    out = self.model(input)  # one forward pass
    loss = self.criterion(out, target)  # calculate loss
    loss_value = loss.item()
    train_losses.append(loss_value)
    loss.backward()  # one backward pass
    self.optimizer.step()  # update the parameters

    batch_iter.set_description(f'Training: (loss {loss_value:.4f})')  # update progressbar

self.training_loss.append(np.mean(train_losses))
self.learning_rate.append(self.optimizer.param_groups[0]['lr'])

batch_iter.close()

И сообщение об ошибке

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

Как я могу это решить?