code: def get_data_loaders(train_batch_size, val_batch_size):
      """Method to setup data loaders: train_loader and val_loader"""
      data_transform = Compose(
          [ToTensor(), Normalize((0.1307,), (0.3081,))]
      )
  
      train_loader = DataLoader(
          MNIST(
              download=True,
              root=".",
              transform=data_transform,
              train=True,
          ),
          batch_size=train_batch_size,
          shuffle=True,
          num_workers=4,
      )
  
      val_loader = DataLoader(
          MNIST(
              download=False,
              root=".",
              transform=data_transform,
              train=False,
          ),
          batch_size=val_batch_size,
          shuffle=False,
          num_workers=4,
      )
      return train_loader, val_loader
  
imports: from argparse import ArgumentParser
  from pathlib import Path
  import torch
  import torch.nn.functional
  from torch import nn
  from torch.optim import SGD
  from torch.optim.lr_scheduler import StepLR
  from torch.utils.data import DataLoader
  from torchvision.datasets import MNIST
  from torchvision.transforms import Compose, Normalize, ToTensor
  from tqdm import tqdm
  from ignite.engine import create_supervised_evaluator, create_supervised_trainer, Events
  from ignite.handlers import Checkpoint, DiskSaver
  from ignite.metrics import Accuracy, Loss, RunningAverage
  from ignite.utils import manual_seed
label_desc: Use image node_id 1 for calling a global function with example usage: get_data_loaders(train_batch_size, val_batch_size) and returns: train_loader,  val_loader