code: def train(
      model, device, train_loader, log_interval, optimizer, epoch
  ):
      model.train()
      for batch_idx, row in enumerate(train_loader):
          data, target = row["image"].to(device), row["digit"].to(
              device
          )
          optimizer.zero_grad()
          output = model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
          if batch_idx % log_interval == 0:
              print(
                  "Train Epoch: {} [{}]\tLoss: {:.6f}".format(
                      epoch, batch_idx * len(data), loss.item()
                  )
              )
  
file_path: petastorm/examples/mnist/pytorch_example.py
imports: from __future__ import division, print_function
  import argparse
  import pyarrow
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from torchvision import transforms
  from examples.mnist import DEFAULT_MNIST_DATA_PATH
  from petastorm import make_reader, TransformSpec
  from petastorm.pytorch import DataLoader
label_desc: Use image node_id 1 for calling a global function with example usage: train(model, device, train_loader, log_interval, optimizer, epoch) without return types