'IoU returns NaN pytorch

So I am are currently trying to solve a segmentation problem in PyTorch and I am making use of IoU as the performance metric. Initially when I ran the model for a certain number of epochs, the IoU score which I obtained was around 0.57. Ever since then, any rerun of the same notebook without making any changes, keeps resulting in NaN and I am not sure why. The intersection between the predicted and ground truth keeps getting displayed as 0.

class IoUEval:
    def __init__(self, nthresh=255):
        self.nthresh = nthresh
        self.thresh = torch.linspace(1./(nthresh + 1), 1. - 1./(nthresh + 1), nthresh).cuda()
        self.EPSILON = np.finfo(np.float).eps

        self.gt_sum = torch.zeros((nthresh,)).cuda()
        self.pred_sum = torch.zeros((nthresh,)).cuda()
        self.num_images = 0
        self.mae = 0
        self.prec = torch.zeros(self.nthresh).cuda()
        self.recall = torch.zeros(self.nthresh).cuda()
        self.iou = 0.


    def add_batch(self, predict, gth):
        for i in range(predict.shape[0]):
            dt = predict[i]; gt = gth[i]
            self.mae += (dt-gt).abs().mean()
#             dt = dt > (dt.mean() * 2)
#             gt = gt > 0.5
            intersect = (dt*gt).sum()
# print(intersect.float())
            iou = intersect.float() / (dt.sum() + gt.sum() - intersect).float()
            if torch.isnan(iou):
                plt.imshow(predict[i].cpu())
                plt.show()
                plt.imshow(gth[i].cpu())
                plt.show()
            else:
                self.iou += iou
            
        self.num_images += predict.shape[0]
        

    def get_metric(self):
        x = self.iou / self.num_images
        y = self.mae / self.num_images
        return x, y
    
    def print_metric(self):
        return self.iou


Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source