-vis 1 RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 256 but got size 3 for tensor number 1 in the list.

[krishnaadithya]

actual shape : img : torch.Size([64, 256, 256, 3]) pred_mask : torch.Size([64, 1, 256, 256]) gt_mask; torch.Size([64, 1, 256, 256])
vis code shape : img : torch.Size([4, 256, 256, 256]) pred_mask : torch.Size([4, 3, 256, 256]) gt_mask; torch.Size([4, 3, 256, 256])