code: def inference(args, ctx):
      # load saved_model
      saved_model = tf.saved_model.load(args.export_dir, tags="serve")
      predict = saved_model.signatures["serving_default"]
  
      # parse function for TFRecords
      def parse_tfr(example_proto):
          feature_def = {
              "label": tf.io.FixedLenFeature(1, tf.int64),
              "image": tf.io.FixedLenFeature(784, tf.int64),
          }
          features = tf.io.parse_single_example(
              serialized=example_proto, features=feature_def
          )
          image = tf.cast(features["image"], dtype=tf.float32) / 255.0
          image = tf.reshape(image, [28, 28, 1])
          label = tf.cast(features["label"], dtype=tf.float32)
          return (image, label)
  
      # define a new tf.data.Dataset (for inferencing)
      ds = tf.data.Dataset.list_files(
          "{}/part-*".format(args.images_labels), shuffle=False
      )
      ds = ds.shard(ctx.num_workers, ctx.worker_num)
      ds = ds.interleave(tf.data.TFRecordDataset)
      ds = ds.map(parse_tfr)
      ds = ds.batch(10)
  
      # create an output file per spark worker for the predictions
      tf.io.gfile.makedirs(args.output)
      output_file = tf.io.gfile.GFile(
          "{}/part-{:05d}".format(args.output, ctx.worker_num), mode="w"
      )
  
      for batch in ds:
          predictions = predict(conv2d_input=batch[0])
          labels = np.reshape(batch[1], -1).astype(np.int)
          preds = np.argmax(predictions["dense_1"], axis=1)
          for x in zip(labels, preds):
              output_file.write("{} {}\n".format(x[0], x[1]))
  
      output_file.close()
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import argparse
  import numpy
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: inference(args, ctx) and returns: image,  label