code: def get_train_dataset(stage, dataset_root):
      if stage == "chairs":
          transforms = OpticalFlowPresetTrain(
              crop_size=(368, 496),
              min_scale=0.1,
              max_scale=1.0,
              do_flip=True,
          )
          return FlyingChairs(
              root=dataset_root, split="train", transforms=transforms
          )
      elif stage == "things":
          transforms = OpticalFlowPresetTrain(
              crop_size=(400, 720),
              min_scale=-0.4,
              max_scale=0.8,
              do_flip=True,
          )
          return FlyingThings3D(
              root=dataset_root,
              split="train",
              pass_name="both",
              transforms=transforms,
          )
      elif stage == "sintel_SKH":  # S + K + H as from paper
          crop_size = (368, 768)
          transforms = OpticalFlowPresetTrain(
              crop_size=crop_size,
              min_scale=-0.2,
              max_scale=0.6,
              do_flip=True,
          )
  
          things_clean = FlyingThings3D(
              root=dataset_root,
              split="train",
              pass_name="clean",
              transforms=transforms,
          )
          sintel = Sintel(
              root=dataset_root,
              split="train",
              pass_name="both",
              transforms=transforms,
          )
  
          kitti_transforms = OpticalFlowPresetTrain(
              crop_size=crop_size,
              min_scale=-0.3,
              max_scale=0.5,
              do_flip=True,
          )
          kitti = KittiFlow(
              root=dataset_root,
              split="train",
              transforms=kitti_transforms,
          )
  
          hd1k_transforms = OpticalFlowPresetTrain(
              crop_size=crop_size,
              min_scale=-0.5,
              max_scale=0.2,
              do_flip=True,
          )
          hd1k = HD1K(
              root=dataset_root,
              split="train",
              transforms=hd1k_transforms,
          )
  
          # As future improvement, we could probably be using a distributed sampler here
          # The distribution is S(.71), T(.135), K(.135), H(.02)
          return 100 * sintel + 200 * kitti + 5 * hd1k + things_clean
      elif stage == "kitti":
          transforms = OpticalFlowPresetTrain(
              # resize and crop params
              crop_size=(288, 960),
              min_scale=-0.2,
              max_scale=0.4,
              stretch_prob=0,
              # flip params
              do_flip=False,
              # jitter params
              brightness=0.3,
              contrast=0.3,
              saturation=0.3,
              hue=0.3 / 3.14,
              asymmetric_jitter_prob=0,
          )
          return KittiFlow(
              root=dataset_root, split="train", transforms=transforms
          )
      else:
          raise ValueError(f"Unknown stage {stage}")
  
file_path: vision/references/optical_flow/train.py
imports: import argparse
  import warnings
  from math import ceil
  from pathlib import Path
  import torch
  import torchvision.models.optical_flow
  import utils
  from presets import OpticalFlowPresetEval, OpticalFlowPresetTrain
  from torchvision.datasets import FlyingChairs, FlyingThings3D, HD1K, KittiFlow, Sintel
label_desc: Use image node_id 1 for calling a global function with example usage: get_train_dataset(stage, dataset_root) and returns: FlyingChairs, FlyingThings3D, unknown, KittiFlow