code: def get_torch_data_loader(
      worker_rank, batch_size, num_workers, transform=None
  ):
      """Get PyTorch DataLoader for the specified training worker.
  
      The input files are split across all workers, and this PyTorch DataLoader
      would only read the portion of files for itself.
      """
      import os
      import numpy as np
      from torchdata.datapipes.iter import IterableWrapper, S3FileLoader
  
      # NOTE: these two variables need to be set to read from S3 successfully.
      os.environ["S3_VERIFY_SSL"] = "0"
      os.environ["AWS_REGION"] = "us-west-2"
  
      def load_image(inputs):
          import io
          from PIL import Image
  
          url, fd = inputs
          data = fd.file_obj.read()
          image = Image.open(io.BytesIO(data))
          image = image.convert("RGB")
          if transform is not None:
              image = transform(
                  pil_to_tensor(image) / 255.0,
              )
          return image
  
      class FileURLDataset:
          """The PyTorch Dataset to split input files URLs among workers."""
  
          def __init__(self, file_urls):
              self._file_urls = file_urls
  
          def __iter__(self):
              worker_info = torch.utils.data.get_worker_info()
              assert worker_info is not None
  
              torch_worker_id = worker_info.id
              return iter(self._file_urls[torch_worker_id])
  
      file_urls = INPUT_FILES_PER_WORKER[worker_rank]
      file_urls = [
          f.tolist() for f in np.array_split(file_urls, num_workers)
      ]
      file_url_dp = IterableWrapper(FileURLDataset(file_urls))
      file_dp = S3FileLoader(file_url_dp)
      image_dp = file_dp.map(load_image)
  
      # NOTE: the separate implementation for using fsspec.
      # Comment out by default. Leave it here as reference.
      #
      # subdir_url_dp = IterableWrapper([root_dir]).list_files_by_fsspec()
      # file_url_dp = subdir_url_dp.list_files_by_fsspec()
      # file_dp = file_url_dp.open_files_by_fsspec(mode="rb")
      # image_dp = file_dp.map(load_image)
  
      data_loader = torch.utils.data.DataLoader(
          image_dp,
          batch_size=batch_size,
          num_workers=num_workers,
      )
      return data_loader
  
imports: import ray
  from ray import train
  from ray.train import DataConfig, ScalingConfig, RunConfig, Checkpoint
  from ray.train.torch import TorchTrainer
  from ray.data.datasource.partitioning import Partitioning
  import tempfile
  import itertools
  import os
  import time
  from benchmark import Benchmark, BenchmarkMetric
  from image_loader_microbenchmark import get_transform, crop_and_flip_image, decode_image_crop_and_flip, center_crop_image
  from image_loader_microbenchmark import get_mosaic_dataloader
  import torch
  import torch.distributed
  from torchvision.models import resnet50
  from torchvision.transforms.functional import pil_to_tensor
  import torch.nn
  import torch.optim
  from dataset_benchmark_util import get_prop_parquet_paths, IMG_S3_ROOT, get_mosaic_epoch_size, IMAGENET_WNID_TO_ID
label_desc: Use image node_id 6 for calling a global function with example usage: get_torch_data_loader(worker_rank, batch_size, num_workers, transform) and returns: data_loader, image, iter