code: def load_dataset(config, split: str):
      """
      Load the segmentation dataset
      """
      image_transforms = transforms.Compose(
          [
              transforms.CenterCrop(size=(224, 224)),
              transforms.ToTensor(),
              transforms.Normalize(
                  mean=[0.485, 0.456, 0.406],
                  std=[0.229, 0.224, 0.225],
                  inplace=True,
              ),
          ]
      )
      target_transforms = transforms.Compose(
          [
              transforms.CenterCrop(size=(224, 224)),
              transforms.ToTensor(),
              transforms.Lambda(
                  lambda x: (x * 255).to(dtype=torch.int64)
              ),
              transforms.Lambda(
                  # TODO - fix this - use an ignore index instead
                  lambda x: torch.where(
                      x != 255, x, torch.tensor(0, dtype=torch.int64)
                  )
              ),
          ]
      )
      data_files, label_files = dataset_catalog.get_data_files(
          split="TRAIN"
          if split == "train"
          else "TEST",  # VISSL split name
          dataset_config=config["DATA"],
      )
      return datasets.VOCSegmentation(
          root=data_files[0],
          image_set=split,
          transform=image_transforms,
          target_transform=target_transforms,
      )
  
file_path: vissl/tools/knn_segmentation.py
imports: import sys
  from typing import Any, List
  import faiss
  import numpy
  import sklearn.metrics
  import torch
  import torch.nn.functional
  import torch.utils.data
  import torchvision.datasets
  import torchvision.transforms
  from tqdm import tqdm
  from vissl.data import dataset_catalog
  from vissl.models import build_model
  from vissl.utils.checkpoint import CheckpointLoader
  from vissl.utils.hydra_config import compose_hydra_configuration, convert_to_attrdict
  from vissl.utils.knn_utils import get_sorted_predictions
label_desc: Use image node_id 1 for calling a global function with example usage: load_dataset(config, split) and returns: datasets