About segmentation loss function

Hi everyone!

I’m doing a project about semantic segmentation. Since I cannot find a good example for segmentation, my project cannot produce good accuracy. The following is some relative codes.
criterion = nn.CrossEntropyLoss().cuda() image, target = image.cuda(), mask.cuda() image, target = Variable(image), Variable(target) output = model(image) _, pred = torch.max(output, dim=1) output = output.permute(0,2,3,1).contiguous() output = output.view(-1, output.size()[-1]) mask_label = target.view(-1) loss = criterion(output, mask_label)
image is BCHW and target is BHW. I wonder whether the codes is correct and anyone knows some function to calculate mean IoU in PyTorch?