code: def train_model(
      net, epochs, ctx, learning_rate, momentum, train_data, val_data
  ):
      # Collect all parameters from net and its children, then initialize them.
      net.initialize(mx.init.Xavier(magnitude=2.24), ctx=ctx)
      # Trainer is for updating parameters with gradient.
      trainer = gluon.Trainer(
          net.collect_params(),
          "sgd",
          {"learning_rate": learning_rate, "momentum": momentum},
      )
      metric = mx.metric.Accuracy()
      loss = gluon.loss.SoftmaxCrossEntropyLoss()
  
      for epoch in range(epochs):
          # reset data iterator and metric at begining of epoch.
          metric.reset()
          for i, (data, label) in enumerate(train_data):
              # Copy data to ctx if necessary
              data = data.as_in_context(ctx)
              label = label.as_in_context(ctx)
              # Start recording computation graph with record() section.
              # Recorded graphs can then be differentiated with backward.
              with autograd.record():
                  output = net(data)
                  L = loss(output, label)
                  L.backward()
              # take a gradient step with batch_size equal to data.shape[0]
              trainer.step(data.shape[0])
              # update metric at last.
              metric.update([label], [output])
  
              if i % 100 == 0 and i > 0:
                  name, acc = metric.get()
                  print(
                      "[Epoch %d Batch %d] Training: %s=%f"
                      % (epoch, i, name, acc)
                  )
  
          name, acc = metric.get()
          print("[Epoch %d] Training: %s=%f" % (epoch, name, acc))
          name, val_acc = test(ctx, net, val_data)
          print("[Epoch %d] Validation: %s=%f" % (epoch, name, val_acc))
  
file_path: sagemaker-python-sdk/tests/data/mxnet_mnist/mnist_gluon.py
imports: import argparse
  import random
  import mxnet
  import numpy
  from mxnet import autograd, gluon
  from mxnet.gluon import nn
label_desc: Use image node_id 3 for calling a global function with example usage: train_model(net, epochs, ctx, learning_rate, momentum, train_data, val_data) without return types