ValueError: Height and Width of image, mask or masks should be equal.disable shapes check by setting is_check_shapes=False of Compose class

PyTorch Live
1

I would like to mention that the shape of image before transformation is → Original image shape: (360, 640, 3) Original mask shape: (360, 640) transformed image shape: (320, 320, 3) transformed mask shape: (320, 320)

and img_trans is defined as → img_trans = A.Compose([A.Resize(a_config.INPUT_IMAGE_HEIGHT, a_config.INPUT_IMAGE_WIDTH)],is_check_shapes=False)

The code snippet is->

def image_mask_transformation(image, mask, img_trans, aug_trans=None, normalize=True,
                              var_min=10, var_max=400):

    
    
    print("Original image shape:", image.shape)
    transformed = img_trans(image=image, mask=mask)
    image = transformed["image"]
    mask = transformed["mask"]
    print("transformed image shape:", image.shape)
    # image, mask still uint 8 in 0, 255 range

But I am getting an error →

ValueError: Caught ValueError in DataLoader worker process 0.

ValueError: Height and Width of image, mask or masks should be equal. You can disable shapes check by setting a parameter is_check_shapes=False of Compose class (do it only if you are sure about your data consistency).

I am not sure whether I am getting the error because initially a 3rd dimension is there for image and not for the mask. Hence I have tried to add the is_check_shapes=False while defining img_trans but I am still getting the error.

Can you suggest whether and where I should add is_check_shapes=False or the error can be solved in any other way?