I assume your mask images contain a specific color code, which would be used to create the corresponding class indices.
If that’s the case, be careful with transforms.Resize((224,224)), as the default interpolation would be a bilinear interpolation, which could corrupt some mask colors.
To resize a mask image you should this use PIL.Image.NEAREST.
Let me know, if this helps.