Only batches of spatial targets supported error

Hello, I was following Fine-Tune a Semantic Segmentation Model with a Custom Dataset with my own dataset. In train.train() it gave me this error “only batches of spatial targets supported (3D tensors) but got targets of size: : [32, 512, 512, 3]”. I followed every steps from tutorial. how can I solve this?