def labels_to_image(img_labels: Tensor, labels_map: Dict[int, Tuple[int, int, int]]) -> Tensor:
        """Function that given an image with labels ids and their pixel intensity mapping, creates an RGB representation
        for visualisation purposes."""
        assert len(img_labels.shape) == 2, img_labels.shape
        H, W = img_labels.shape
        out = torch.empty(3, H, W, dtype=torch.uint8)
        for label_id, label_val in labels_map.items():
            mask = img_labels == label_id
            for i in range(3):
                out[i].masked_fill_(mask, label_val[i])
        return out
Function that given an image with labels ids and their pixel intensity mapping, creates an RGB representation
        for visualisation purposes.
src/flash/image/segmentation/output.py