text: import torch
  
  EPS = 1e-6
  
  def iou_single(a, b, mask, n_class):
      valid = mask == 1
      a = a[valid]
      b = b[valid]
      miou = []
      for i in range(n_class):
          inter = ((a == i) & (b == i)).float()
          union = ((a == i) | (b == i)).float()
  
          miou.append(torch.sum(inter) / (torch.sum(union) + EPS))
      miou = sum(miou) / len(miou)
      return miou
  
  def iou(a, b, mask, n_class=2, reduce=True):
      batch_size = a.size(0)
  
      a = a.view(batch_size, -1)
      b = b.view(batch_size, -1)
      mask = mask.view(batch_size, -1)
  
      iou = a.new_zeros((batch_size,), dtype=torch.float32)
      for i in range(batch_size):
          iou[i] = iou_single(a[i], b[i], mask[i], n_class)
  
      if reduce:
          iou = torch.mean(iou)
      return iou