code: def convert_raw_predictions(
      batch,
      raw_preds: torch.Tensor,
      records: Sequence[BaseRecord],
      keep_images: bool = False,
  ) -> List[Prediction]:
      mask_preds = raw_preds.argmax(dim=1)
  
      if len(batch) > 1:
          tensor_images, tensor_gts = batch
      else:
          tensor_images = batch[0]
          tensor_gts = [None] * len(records)
  
      # PL tensors are only 3 dimensions
      if tensor_images.dim() == 3:
          tensor_images = tensor_images[None, :]
  
      preds = []
  
      if len(tensor_images) != len(records):
          raise (
              Exception(
                  "Tensor images and records should have the same length"
              )
          )
  
      for record, tensor_image, mask_pred, tensor_gt in zip(
          records, tensor_images, mask_preds, tensor_gts
      ):
          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.cpu().numpy()[None])
          )
  
          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.cpu().numpy())
                  )
  
          if keep_images:
              record.set_img(tensor_to_image(tensor_image))
  
          preds.append(Prediction(pred=pred, ground_truth=record))
  
      return preds
  
file_path: icevision/icevision/models/fastai/unet/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.fastai.unet.dataloaders import *
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