code: def _collate_fn(examples, with_prior_preservation=False):
      if with_prior_preservation:
          prompts = [
              [example["instance_prompt"], example["reg_prompt"]]
              for example in examples
          ]
          images = [
              example["instance_images"] for example in examples
          ] + [example["reg_images"] for example in examples]
      else:
          prompts = [
              [example["instance_prompt"]] for example in examples
          ]
          images = [example["instance_images"] for example in examples]
  
      images = torch.stack(images)
      images = images.to(memory_format=torch.contiguous_format).float()
  
      return prompts, images
  
imports: from functools import partial
  from typing import Any, Optional
  import torch
  from omegaconf import DictConfig
  from pytorch_lightning import Trainer
  from torch._inductor import config
  from nemo.collections.multimodal.data.dreambooth.dreambooth_dataset import DreamBoothDataset
  from nemo.collections.multimodal.modules.stable_diffusion.distributions.distributions import DiagonalGaussianDistribution
  from nemo.collections.multimodal.parts.utils import randn_like
  from nemo.collections.nlp.data.language_modeling.megatron.data_samplers import MegatronPretrainingRandomSampler
  from nemo.collections.nlp.models.language_modeling.megatron_base_model import MegatronBaseModel
  from nemo.collections.nlp.modules.common.megatron.module import Float16Module
  from nemo.collections.nlp.parts.utils_funcs import get_last_rank
  from nemo.core.classes.common import Serialization
  from nemo.utils import logging
label_desc: Use image node_id 2 for calling a global function with example usage: _collate_fn(examples, with_prior_preservation) and returns: prompts,  images