def METHOD_NAME(
    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_training_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 current_host == scheduler_host(hosts):
        save(model_dir, mlp_model)
        assert_can_track_sagemaker_experiments()