code: def flatten_probs(probs, labels, ignore_index=None):
      """Flattens predictions in the batch."""
      if probs.dim() == 3:
          # assumes output of a sigmoid layer
          B, H, W = probs.size()
          probs = probs.view(B, 1, H, W)
      B, C, H, W = probs.size()
      probs = (
          probs.permute(0, 2, 3, 1).contiguous().view(-1, C)
      )  # B*H*W, C=P,C
      labels = labels.view(-1)
      if ignore_index is None:
          return probs, labels
      valid = labels != ignore_index
      vprobs = probs[valid.nonzero().squeeze()]
      vlabels = labels[valid]
      return vprobs, vlabels
  
file_path: lavis/lavis/common/annotator/uniformer/mmseg/models/losses/lovasz_loss.py
imports: import annotator.uniformer.mmcv
  import torch
  import torch.nn
  import torch.nn.functional
  from ..builder import LOSSES
  from .utils import get_class_weight, weight_reduce_loss
label_desc: Use image node_id 3 for calling a global function with example usage: flatten_probs(probs, labels, ignore_index) and returns: vprobs,  vlabels, probs,  labels