code: def load_data(
      *,
      data_dir,
      batch_size,
      image_size,
      class_cond=False,
      deterministic=False,
  ):
      """
      For a dataset, create a generator over (images, kwargs) pairs.
  
      Each images is an NCHW float tensor, and the kwargs dict contains zero or
      more keys, each of which map to a batched Tensor of their own.
      The kwargs dict can be used for class labels, in which case the key is "y"
      and the values are integer tensors of class labels.
  
      :param data_dir: a dataset directory.
      :param batch_size: the batch size of each returned pair.
      :param image_size: the size to which images are resized.
      :param class_cond: if True, include a "y" key in returned dicts for class
                         label. If classes are not available and this is true, an
                         exception will be raised.
      :param deterministic: if True, yield results in a deterministic order.
      """
      if not data_dir:
          raise ValueError("unspecified data directory")
      all_files = _list_image_files_recursively(data_dir)
      classes = None
      if class_cond:
          # Assume classes are the first part of the filename,
          # before an underscore.
          class_names = [
              bf.basename(path).split("_")[0] for path in all_files
          ]
          sorted_classes = {
              x: i for i, x in enumerate(sorted(set(class_names)))
          }
          classes = [sorted_classes[x] for x in class_names]
      dataset = ImageDataset(
          image_size,
          all_files,
          classes=classes,
          shard=MPI.COMM_WORLD.Get_rank(),
          num_shards=MPI.COMM_WORLD.Get_size(),
      )
      if deterministic:
          loader = DataLoader(
              dataset,
              batch_size=batch_size,
              shuffle=False,
              num_workers=1,
              drop_last=True,
          )
      else:
          loader = DataLoader(
              dataset,
              batch_size=batch_size,
              shuffle=True,
              num_workers=1,
              drop_last=True,
          )
      while True:
          yield from loader
  
file_path: improved-diffusion/improved_diffusion/image_datasets.py
imports: from PIL import Image
  import blobfile
  from mpi4py import MPI
  import numpy
  from torch.utils.data import DataLoader, Dataset
label_desc: Use image node_id 1 for calling a global function with example usage: load_data() without return types