Thanks for the explanation.
Actually, I am performing a volumetric segmentation using a 2D slice wise approach. The image to be fed into the model has a spatial dimension of [512, 512] and has there one-hot encoded labels of the same dimension.
import torch
from torchvision import transforms

image = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float64)
labelA = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float64)
labelB = torch.randint(low=0, high=2, size=(3, 3), dtype=torch.float64)

# labelC is created by me which is the backgound of labelA and label B i.e it zeros where either A or B is 1
labelC = torch.zeros((labelB.shape), dtype=torch.long)

for label in [labelA, labelB]:
    labelC |= label.long()
labelC = labelC^1

labels = torch.cat([labelC.unsqueeze_(0)] + [x.unsqueeze_(0) for x in [labelA, labelB]] , dim = 0)
print(labelA, labelB.long(), labelC)
print(labels.shape)


# print(labels.shape)
image_trans = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize((3, 3)),
            transforms.ToTensor(),
            transforms.Normalize([0.5], [0.5])])

gt_trans = transforms.Compose([
        transforms.Resize((3, 3)),
])
transformed_image =  image_trans(image)
transformed_labels = gt_trans(labels)
print(transformed_image.shape, transformed_labels.shape)
print(transformed_image.dtype, transformed_labels.dtype)