code: def flatten_probas(probas, labels, ignore=None):
      
      B, C, H, W = probas.size()
      probas = probas.permute(0, 2, 3, 1).contiguous().view(-1, C)  # B * H * W, C = P, C
      labels = labels.view(-1)
      if ignore is None:
          return probas, labels
      valid = (labels != ignore)
      vprobas = probas[valid.nonzero().squeeze()]
      vlabels = labels[valid]
      return vprobas, vlabels
docs: Flattens predictions in the batch