code: def load_data(data_dir):
      transform = transforms.Compose(
          [
              transforms.ToTensor(),
              transforms.Normalize(
                  (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
              ),
          ]
      )
  
      # We add FileLock here because multiple workers will want to
      # download data, and this may cause overwrites since
      # DataLoader is not threadsafe.
      with FileLock(os.path.expanduser("~/.data.lock")):
          trainset = torchvision.datasets.CIFAR10(
              root=data_dir,
              train=True,
              download=True,
              transform=transform,
          )
  
          testset = torchvision.datasets.CIFAR10(
              root=data_dir,
              train=False,
              download=True,
              transform=transform,
          )
  
      return trainset, testset
  
imports: import numpy
  import os
  import tempfile
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from filelock import FileLock
  from torch.utils.data import random_split
  import torchvision
  import torchvision.transforms
  from typing import Dict
  import ray
  from ray import train, tune
  from ray.train import Checkpoint
  from ray.tune.schedulers import ASHAScheduler
label_desc: Use image node_id 1 for calling a global function with example usage: load_data(data_dir) and returns: trainset,  testset