The output shape is not correct

[Paranoidv]

When I test picture from VOC2017,I get the wrong message 🔢 Ucat
Traceback (most recent call last):
File "tester.py", line 46, in
output = Colorize()(output)
File "/home/xiangjie/Documents/master_thesis/segmentation/pytorch-seg/transform.py", line 150, in call
color_image[0][mask] = self.cmap[label][0]
RuntimeError: inconsistent tensor size, expected tensor [256 x 256] and mask [256] to have the same number of elements, but got 65536 and 256 elements respectively at /pytorch/torch/lib/TH/generic/THTensorMath.c:138

I try to fix it.But did not work ,so ,someone can tell me how to fix this bug????
Thank you!!!!