code: def collate(batch, dtype):
      """Build Torch training batch.
  
      B = batch size
      (C, W, H) = (channels, width, height)
      L = max length in tokens of the text guidance input
  
      Input batch schema (see `get_train_dataset` on how this was setup):
          instance_images: (B, C, W, H)
          class_images: (B, C, W, H)
          instance_prompt_ids: (B, L)
          class_prompt_ids: (B, L)
  
      Output batch schema:
          images: (2 * B, C, W, H)
              All instance images in the batch come before the class images:
              [instance_images[0], ..., instance_images[B-1], class_images[0], ...]
          prompt_ids: (2 * B, L)
              Prompt IDs are ordered the same way as the images.
  
      During training, a batch will be chunked into 2 sub-batches for
      prior preserving loss calculation.
      """
  
      images = torch.cat(
          [batch["instance_image"], batch["class_image"]], dim=0
      )
      images = images.to(memory_format=torch.contiguous_format).to(
          dtype
      )
  
      batch_size = len(batch["instance_prompt_ids"])
  
      prompt_ids = torch.cat(
          [batch["instance_prompt_ids"], batch["class_prompt_ids"]],
          dim=0,
      ).reshape(batch_size * 2, -1)
  
      return {
          "images": images,
          "prompt_ids": prompt_ids,  # token ids should stay int.
      }
  
file_path: ray/release/workspace_templates/05_dreambooth_finetuning/dreambooth/dataset.py
imports: from typing import Dict
  import numpy
  import pandas
  import torch
  from ray.data import read_images
  from torchvision import transforms
  from transformers import AutoTokenizer
label_desc: Use image node_id 2 for calling a global function with example usage: collate(batch, dtype) and returns: dict