code: def train(
      batch_size,
      epochs,
      learning_rate,
      num_gpus,
      training_channel,
      testing_channel,
      hosts,
      current_host,
      model_dir,
  ):
      (train_labels, train_images) = load_data(training_channel)
      (test_labels, test_images) = load_data(testing_channel)
  
      # Data parallel training - shard the data so each host
      # only trains on a subset of the total data.
      shard_size = len(train_images) // len(hosts)
      for i, host in enumerate(hosts):
          if host == current_host:
              start = shard_size * i
              end = start + shard_size
              break
  
      train_iter = mx.io.NDArrayIter(
          train_images[start:end],
          train_labels[start:end],
          batch_size,
          shuffle=True,
      )
      val_iter = mx.io.NDArrayIter(test_images, test_labels, batch_size)
  
      logging.getLogger().setLevel(logging.DEBUG)
  
      kvstore = "local" if len(hosts) == 1 else "dist_sync"
  
      mlp_model = mx.mod.Module(
          symbol=build_graph(), context=get_train_context(num_gpus)
      )
      mlp_model.fit(
          train_iter,
          eval_data=val_iter,
          kvstore=kvstore,
          optimizer="sgd",
          optimizer_params={"learning_rate": learning_rate},
          eval_metric="acc",
          batch_end_callback=mx.callback.Speedometer(batch_size, 100),
          num_epoch=epochs,
      )
  
      if len(hosts) == 1 or current_host == scheduler_host(hosts):
          save(model_dir, mlp_model)
  
file_path: sagemaker-python-sdk/tests/data/mxnet_mnist/mnist_neo.py
imports: from __future__ import absolute_import
  import argparse
  import gzip
  import json
  import logging
  import os
  import struct
  import mxnet
  import numpy
label_desc: Use image node_id 5 for calling a global function with example usage: train(batch_size, epochs, learning_rate, num_gpus, training_channel, testing_channel, hosts, current_host, model_dir) without return types