code: def train(
      model,
      graph,
      dataloader,
      optimizer,
      scheduler,
      normalizer,
      loss_fn,
      device,
      args,
  ):
      total_loss = []
      graph = graph.to(device)
      model.train()
      batch_size = args.batch_size
      for i, (x, y) in enumerate(dataloader):
          optimizer.zero_grad()
          # Padding: Since the diffusion graph is precmputed we need to pad the batch so that
          # each batch have same batch size
          if x.shape[0] != batch_size:
              x_buff = torch.zeros(
                  batch_size, x.shape[1], x.shape[2], x.shape[3]
              )
              y_buff = torch.zeros(
                  batch_size, x.shape[1], x.shape[2], x.shape[3]
              )
              x_buff[: x.shape[0], :, :, :] = x
              x_buff[x.shape[0] :, :, :, :] = x[-1].repeat(
                  batch_size - x.shape[0], 1, 1, 1
              )
              y_buff[: x.shape[0], :, :, :] = y
              y_buff[x.shape[0] :, :, :, :] = y[-1].repeat(
                  batch_size - x.shape[0], 1, 1, 1
              )
              x = x_buff
              y = y_buff
          # Permute the dimension for shaping
          x = x.permute(1, 0, 2, 3)
          y = y.permute(1, 0, 2, 3)
  
          x_norm = (
              normalizer.normalize(x)
              .reshape(x.shape[0], -1, x.shape[3])
              .float()
              .to(device)
          )
          y_norm = (
              normalizer.normalize(y)
              .reshape(x.shape[0], -1, x.shape[3])
              .float()
              .to(device)
          )
          y = y.reshape(y.shape[0], -1, y.shape[3]).float().to(device)
  
          batch_graph = dgl.batch([graph] * batch_size)
          output = model(
              batch_graph, x_norm, y_norm, batch_cnt[0], device
          )
          # Denormalization for loss compute
          y_pred = normalizer.denormalize(output)
          loss = loss_fn(y_pred, y)
          loss.backward()
          nn.utils.clip_grad_norm_(
              model.parameters(), args.max_grad_norm
          )
          optimizer.step()
          if get_learning_rate(optimizer) > args.minimum_lr:
              scheduler.step()
          total_loss.append(float(loss))
          batch_cnt[0] += 1
          print("\rBatch: ", i, end="")
      return np.mean(total_loss)
  
file_path: dgl/examples/pytorch/dtgrnn/train.py
imports: import argparse
  from functools import partial
  import dgl
  import numpy
  import torch
  import torch.nn
  from dataloading import METR_LAGraphDataset, METR_LATestDataset, METR_LATrainDataset, METR_LAValidDataset, PEMS_BAYGraphDataset, PEMS_BAYTestDataset, PEMS_BAYTrainDataset, PEMS_BAYValidDataset
  from dcrnn import DiffConv
  from gaan import GatedGAT
  from model import GraphRNN
  from torch.utils.data import DataLoader
  from utils import get_learning_rate, masked_mae_loss, NormalizationLayer
label_desc: Use image node_id 1 for calling a global function with example usage: train(model, graph, dataloader, optimizer, scheduler, normalizer, loss_fn, device, args) and returns: np