code: def train(data_loader, steps=100, lr=0.0005, momentum=0.5):
      model = Net()
      model.train()
      optimizer = optim.SGD(
          model.parameters(), lr=lr, momentum=momentum
      )
      loss_hist = []
      for batch_idx, batch in enumerate(data_loader):
          if batch_idx > steps:
              break
          data, target = Variable(batch["features"]), Variable(
              batch["label"]
          )
          optimizer.zero_grad()
          output = model(data)
          loss = F.nll_loss(output, target)
          loss.backward()
          optimizer.step()
          if batch_idx % 10 == 0:
              logging.info(
                  "[%d/%d]\tLoss: %.6f",
                  batch_idx,
                  steps,
                  loss.data.item(),
              )
              loss_hist.append(loss.data.item())
      return model
  
file_path: petastorm/examples/spark_dataset_converter/pytorch_converter_example.py
imports: from __future__ import division
  import logging
  import torch
  import torch.nn
  import torch.nn.functional
  import torch.optim
  from pyspark.sql import SparkSession
  from torch.autograd import Variable
  from examples.spark_dataset_converter.utils import get_mnist_dir
  from petastorm.spark import SparkDatasetConverter, make_spark_converter
label_desc: Use image node_id 1 for calling a global function with example usage: train(data_loader, steps, lr, momentum) and returns: model