code: def load_dataset(dataset_dir, batch_size=4):
      # Data augmentation and normalization for training
      # Just normalization for validation
      data_transforms = {
          "train": transforms.Compose(
              [
                  transforms.RandomResizedCrop(224),
                  transforms.RandomHorizontalFlip(),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
                  ),
              ]
          ),
          "val": transforms.Compose(
              [
                  transforms.Resize(256),
                  transforms.CenterCrop(224),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
                  ),
              ]
          ),
      }
  
      image_datasets = {
          x: datasets.ImageFolder(
              os.path.join(dataset_dir, x), data_transforms[x]
          )
          for x in ["train", "val"]
      }
      dataset_sizes = {
          x: len(image_datasets[x]) for x in ["train", "val"]
      }
      class_names = image_datasets["train"].classes
  
      dataloaders = {
          x: torch.utils.data.DataLoader(
              image_datasets[x],
              batch_size=batch_size,
              shuffle=True,
              num_workers=1,
          )
          for x in ["train", "val"]
      }
      train_loader = dataloaders["train"]
      test_loader = dataloaders["val"]
      return train_loader, test_loader, class_names
  
file_path: BigDL/python/orca/tutorial/pytorch/transfer_learning/train.py
imports: import os
  import numpy
  import matplotlib.pyplot
  import torch
  import torch.nn
  import torch.optim
  from torch.optim import lr_scheduler
  import torchvision
  from torchvision import datasets, transforms
  from bigdl.orca import init_orca_context, stop_orca_context
  from bigdl.orca.learn.pytorch import Estimator
  from bigdl.orca.learn.metrics import Accuracy
label_desc: Use image node_id 1 for calling a global function with example usage: load_dataset(dataset_dir, batch_size) and returns: train_loader,  test_loader,  class_names