code: def convert_raw_predictions(
      batch,
      raw_preds: torch.Tensor,
      records: Sequence[BaseRecord],
      keep_images: bool = False,
  ) -> List[Prediction]:
      # Retrieve ground_truth masks from batch, if present (at training time)
      if "gt_semantic_seg" in batch:
          tensor_gts = (
              batch["gt_semantic_seg"].squeeze().chunk(8, dim=0)
          )
      else:
          tensor_gts = [None]
  
      # Images for inference come in as lists
      if isinstance(batch["img"], list):
          tensor_imgs = [
              x.chunk(batch["img"][0].shape[0], dim=0)
              for x in batch["img"]
          ]
          tensor_imgs = list(
              itertools.chain.from_iterable(tensor_imgs)
          )  # Ensuring that all chunks are unrolled into the list
      else:
          # At training time, images are in a single tensor
          tensor_imgs = batch["img"].chunk(batch["img"].shape[0], dim=0)
  
      if len(tensor_imgs) != len(records):
          raise (
              Exception(
                  "Tensor images and records should have the same length"
              )
          )
  
      # TODO: This can be optimized by using looping through first tensor dims behaviour instead of chunking tensors?
      preds = []
      for record, tensor_gt, mask_pred, tensor_image in zip(
          records, cycle(tensor_gts), raw_preds, tensor_imgs
      ):
          pred = BaseRecord(
              (
                  ImageRecordComponent(),
                  SemanticMaskRecordComponent(),
                  ClassMapRecordComponent(task=tasks.segmentation),
              )
          )
  
          pred.segmentation.set_class_map(record.segmentation.class_map)
          pred.segmentation.set_mask_array(MaskArray(mask_pred))
  
          if tensor_gt is not None:
              # This is used at train time to have mask available for metric computation
              if torch.is_tensor(tensor_gt):
                  record.segmentation.set_mask_array(
                      MaskArray(tensor_gt.squeeze().cpu().numpy())
                  )
  
          if keep_images:
              record.set_img(tensor_to_image(tensor_image.squeeze()))
  
              # We load the record to ensure mask data is available
              if len(record.segmentation.masks):
                  (h, w, _) = record.img.shape
  
                  record.segmentation.set_mask_array(
                      record.segmentation.masks[0].to_mask(
                          h=h, w=w, pad_dim=False
                      )
                  )
  
          preds.append(Prediction(pred=pred, ground_truth=record))
  
      return preds
  
file_path: icevision/icevision/models/mmseg/common/prediction.py
imports: from icevision.imports import *
  from icevision.utils import *
  from icevision.core import *
  from icevision.data import *
  from icevision.models.utils import _predict_from_dl
  from icevision.models.mmseg.common.utils import *
  from icevision.models.mmseg.common.dataloaders import *
  from icevision.models.interpretation import _move_to_device
  import itertools
label_desc: Use image node_id 2 for calling a global function with example usage: convert_raw_predictions(batch, raw_preds, records, keep_images) and returns: preds