code: def main(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")
  
file_path: sagemaker-python-sdk/tests/data/tensorflow_mnist/mnist_v2.py
imports: import argparse
  import numpy
  import os
  import json
  import tensorflow
  from tensorflow.keras.layers import Conv2D, BatchNormalization, Dense, Flatten
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types