Binary Mask Semantic Segmentation Issue

[aylinSyntonym]

I got following error with custom dataset which has binary mask. Can you help me?

ValueError: Input and output must have the same number of spatial dimensions, but got input with with spatial dimensions of [128, 128] and output size of torch.Size([512, 512, 3]). Please provide input tensor in (N, C, d1, d2, ...,dK) format and output size in (o1, o2, ...,oK) format.

I solve problem with ADE20K dataset annotation preprocess and decrease the dimension of annotation masks to 2D!