code: def build_train_valid_datasets(model_cfg, data_path, image_size=224):
      if model_cfg.vision_pretraining_type == "classify":
          train_transform = ClassificationTransform(
              model_cfg, image_size
          )
          val_transform = ClassificationTransform(
              model_cfg, image_size, train=False
          )
      elif model_cfg.vision_pretraining_type == "inpaint":
          train_transform = InpaintingTransform(
              model_cfg, image_size, train=False
          )
          val_transform = InpaintingTransform(
              model_cfg, image_size, train=False
          )
      elif model_cfg.vision_pretraining_type == "dino":
          train_transform = DinoTransform(
              model_cfg, image_size, train=True
          )
          val_transform = ClassificationTransform(
              model_cfg, image_size, train=False
          )
      else:
          raise Exception(
              "{} vit pretraining type is not supported.".format(
                  model_cfg.vit_pretraining_type
              )
          )
  
      # training dataset
      train_data_path = (
          data_path[0] if len(data_path) <= 2 else data_path[2]
      )
      train_data = ImageFolder(
          root=train_data_path,
          transform=train_transform,
          classes_fraction=model_cfg.classes_fraction,
          data_per_class_fraction=model_cfg.data_per_class_fraction,
      )
      train_data = RandomSeedDataset(train_data)
  
      # validation dataset
      val_data_path = data_path[1]
      val_data = ImageFolder(
          root=val_data_path, transform=val_transform
      )
      val_data = RandomSeedDataset(val_data)
  
      return train_data, val_data
  
imports: import random
  import numpy
  import torch
  from PIL import Image, ImageFilter, ImageOps
  from torch.utils.data import Dataset
  from nemo.collections.multimodal.data.common.data_samplers import SharedEpoch
  from nemo.collections.vision.data.megatron.autoaugment import ImageNetPolicy
  from nemo.collections.vision.data.megatron.image_folder import ImageFolder
label_desc: Use image node_id 2 for calling a global function with example usage: build_train_valid_datasets(model_cfg, data_path, image_size) and returns: train_data,  val_data