code: def get_dataset(is_train, args):
      image_set = "train" if is_train else "val"
      num_classes, mode = {
          "coco": (91, "instances"),
          "coco_kp": (2, "person_keypoints"),
      }[args.dataset]
      with_masks = "mask" in args.model
      ds = get_coco(
          root=args.data_path,
          image_set=image_set,
          transforms=get_transform(is_train, args),
          mode=mode,
          use_v2=args.use_v2,
          with_masks=with_masks,
      )
      return ds, num_classes
  
file_path: vision/references/detection/train.py
imports: import datetime
  import os
  import time
  import presets
  import torch
  import torch.utils.data
  import torchvision
  import torchvision.models.detection
  import torchvision.models.detection.mask_rcnn
  import utils
  from coco_utils import get_coco
  from engine import evaluate, train_one_epoch
  from group_by_aspect_ratio import create_aspect_ratio_groups, GroupedBatchSampler
  from torchvision.transforms import InterpolationMode
  from transforms import SimpleCopyPaste
label_desc: Use image node_id 2 for calling a global function with example usage: get_dataset(is_train, args) and returns: ds,  num_classes