Monai : RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 7 but got size 8 for tensor number 1 in the list

Question

I am using Monai for the 3D Multilabel segmentation task. My input image size is 512x496x49 and my label size is 512x496x49. An Image can have 3 labels in one image. With transform, I have converted the image in size 1x512x512x49 and Label in 3x512x512x49

My Transform

# Setting tranform for train and test data
a_min=6732
a_max=18732

train_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
      ScaleIntensityRanged(keys='image', a_min=a_min, a_max=a_max, b_min=0.0, b_max=1.0, clip=False),
      Orientationd(keys=["image", "label"], axcodes="RAS"),
      # Spacingd(keys=["image", "label"], pixdim=(
      #     1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=0),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=1),
      RandFlipd(keys=["image", "label"], prob=0.5, spatial_axis=2),
      CropForegroundd(keys=["image", "label"], source_key="image"),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      SpatialPadd(keys=['image', 'label'], spatial_size= [512, 512, 49]),# it will result in 512x512x49
      EnsureTyped(keys=["image", "label"]),
    ]
)
val_transform = Compose(
    [
      LoadImaged(keys=["image", "label"]),
      EnsureChannelFirstd(keys="image"),
      ConvertToMultiChannelBasedOnBratsClassesd(keys="label"),
      ScaleIntensityRanged(keys='image', a_min=a_min, a_max=a_max, b_min=0.0, b_max=1.0, clip=False),
      Orientationd(keys=["image", "label"], axcodes="RAS"),
      # Spacingd(keys=["image", "label"], pixdim=(
      #     1.5, 1.5, 2.0), mode=("bilinear", "nearest")),
      CropForegroundd(keys=["image", "label"], source_key="image"),
      NormalizeIntensityd(keys="image", nonzero=True, channel_wise=True),
      SpatialPadd(keys=['image', 'label'], spatial_size= [512, 512, 49]),# it will result in 512x512x49
      EnsureTyped(keys=["image", "label"]),
    ]
)

Dataloader for training and val

train_ds = CacheDataset(data=train_files, transform=train_transform,cache_rate=1.0, num_workers=4)
train_loader = DataLoader(train_ds, batch_size=2, shuffle=True, num_workers=4,collate_fn=pad_list_data_collate)

val_ds = CacheDataset(data=val_files, transform=val_transform, cache_rate=1.0, num_workers=4)
val_loader = DataLoader(val_ds, batch_size=1, num_workers=4)

3D U-Net Network from Monai

# standard PyTorch program style: create UNet, DiceLoss and Adam optimizer

device = torch.device("cuda:0")
model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=4,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.BATCH,
).to(device)
loss_function = DiceLoss(to_onehot_y=True, sigmoid=True)
optimizer = torch.optim.Adam(model.parameters(), 1e-4)
dice_metric = DiceMetric(include_background=True, reduction="mean")

Training

max_epochs = 5
val_interval = 2
best_metric = -1
best_metric_epoch = -1
epoch_loss_values = []
metric_values = []
post_pred = Compose([EnsureType(), AsDiscrete(argmax=True, to_onehot=4)])
post_label = Compose([EnsureType(), AsDiscrete(to_onehot=4)])

for epoch in range(max_epochs):
    print("-" * 10)
    print(f"epoch {epoch + 1}/{max_epochs}")
    model.train()
    epoch_loss = 0
    step = 0
    for batch_data in train_loader:
        step += 1
        inputs, labels = (
            batch_data["image"].to(device),
            batch_data["label"].to(device),
        )
        optimizer.zero_grad()
        print("Size of inputs :", inputs.shape)
        print("Size of inputs[0] :", inputs[0].shape)
        # print("Size of inputs[1] :", inputs[1].shape)
        # print("printing of inputs :", inputs)
        outputs = model(inputs)
        loss = loss_function(outputs, labels)
        loss.backward()
        optimizer.step()
        epoch_loss += loss.item()
        print(
            f"{step}/{len(train_ds) // train_loader.batch_size}, "
            f"train_loss: {loss.item():.4f}")
    epoch_loss /= step
    epoch_loss_values.append(epoch_loss)
    print(f"epoch {epoch + 1} average loss: {epoch_loss:.4f}")

    if (epoch + 1) % val_interval == 0:
        model.eval()
        with torch.no_grad():
            for val_data in val_loader:
                val_inputs, val_labels = (
                    val_data["image"].to(device),
                    val_data["label"].to(device),
                )
                roi_size = (160, 160, 160)
                sw_batch_size = 4
                val_outputs = sliding_window_inference(
                    val_inputs, roi_size, sw_batch_size, model)
                val_outputs = [post_pred(i) for i in decollate_batch(val_outputs)]
                val_labels = [post_label(i) for i in decollate_batch(val_labels)]
                # compute metric for current iteration
                dice_metric(y_pred=val_outputs, y=val_labels)

            # aggregate the final mean dice result
            metric = dice_metric.aggregate().item()
            # reset the status for next validation round
            dice_metric.reset()

            metric_values.append(metric)
            if metric > best_metric:
                best_metric = metric
                best_metric_epoch = epoch + 1
                torch.save(model.state_dict(), os.path.join(
                    root_dir, "best_metric_model.pth"))
                print("saved new best metric model")
            print(
                f"current epoch: {epoch + 1} current mean dice: {metric:.4f}"
                f"\nbest mean dice: {best_metric:.4f} "
                f"at epoch: {best_metric_epoch}"
            )

While training I am getting this error

RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 7 but got size 8 for tensor number 1 in the list.

I followed the 3D Segmentation Monai tutorial but this was only for 2 classes (including background) therefore I followed the discussion at https://github.com/Project-MONAI/MONAI/issues/415 but even though I changed what was recommended in this discussion still am getting errors while training.