code: def get_dataflow(config):
      # - Get train/test datasets
      with idist.one_rank_first(local=True):
          train_dataset, test_dataset = utils.get_train_test_datasets(
              config["data_path"]
          )
  
      # Setup data loader also adapted to distributed config: nccl, gloo, xla-tpu
      train_loader = idist.auto_dataloader(
          train_dataset,
          batch_size=config["batch_size"],
          num_workers=config["num_workers"],
          shuffle=True,
          drop_last=True,
      )
  
      test_loader = idist.auto_dataloader(
          test_dataset,
          batch_size=2 * config["batch_size"],
          num_workers=config["num_workers"],
          shuffle=False,
      )
      return train_loader, test_loader
  
imports: from datetime import datetime
  from pathlib import Path
  from typing import Any, Optional
  import fire
  import torch
  import torch.nn
  import torch.optim
  import utils
  from torch.cuda.amp import autocast, GradScaler
  import ignite
  import ignite.distributed
  from ignite.contrib.engines import common
  from ignite.contrib.handlers import PiecewiseLinear
  from ignite.engine import Engine, Events
  from ignite.handlers import Checkpoint, DiskSaver, global_step_from_engine
  from ignite.metrics import Accuracy, Loss
  from ignite.utils import manual_seed, setup_logger
label_desc: Use image node_id 3 for calling a global function with example usage: get_dataflow(config) and returns: train_loader,  test_loader