def METHOD_NAME(args):
    if args.model_dir.startswith("s3://"):
        os.environ["S3_REGION"] = "us-west-2"
        os.environ["TF_CPP_MIN_LOG_LEVEL"] = "1"
        os.environ["S3_USE_HTTPS"] = "1"
    net = LeNet()
    net.build(input_shape=(None, 28, 28, 1))
    optimizer = tf.keras.optimizers.Adam()
    train_loss = tf.keras.metrics.Mean(name="train_loss")
    train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="train_accuracy")
    test_loss = tf.keras.metrics.Mean(name="test_loss")
    test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(name="test_accuracy")
    ckpt = tf.train.Checkpoint(optimizer=optimizer, model=net)
    ckpt_manager = tf.train.CheckpointManager(
        ckpt, args.model_dir, max_to_keep=11, checkpoint_name="model.ckpt"
    )
    dtrain, deval = load_data(args.train)
    num_epochs = args.epochs
    for i in range(num_epochs):
        for x, y in dtrain:
            train_step(x, y, net, optimizer, train_loss, train_accuracy)
        for x, y in deval:
            eval_step(x, y, net, test_loss, test_accuracy)
        print(
            f"Epoch {i+1}",
            f"Train Loss: {train_loss.result()}",
            f"Train Accuracy: {train_accuracy.result()}",
            f"Test Loss: {test_loss.result()}",
            f"Test Accuracy: {test_accuracy.result()}",
        )
        if args.current_host == args.hosts[0]:
            ckpt_manager.save()
            net.save("/opt/ml/model/1")