code: def setup(self, config):
      # IMPORTANT: See the above note.
      import tensorflow as tf
  
      # Use FileLock to avoid race conditions.
      with FileLock(os.path.expanduser("~/.tune.lock")):
          (x_train, y_train), (x_test, y_test) = load_data()
      x_train, x_test = x_train / 255.0, x_test / 255.0
  
      # Add a channels dimension
      x_train = x_train[..., tf.newaxis]
      x_test = x_test[..., tf.newaxis]
      self.train_ds = tf.data.Dataset.from_tensor_slices(
          (x_train, y_train)
      )
      self.train_ds = self.train_ds.shuffle(10000).batch(
          config.get("batch", 32)
      )
  
      self.test_ds = tf.data.Dataset.from_tensor_slices(
          (x_test, y_test)
      ).batch(32)
  
      self.model = MyModel(hiddens=config.get("hiddens", 128))
      self.loss_object = tf.keras.losses.SparseCategoricalCrossentropy()
      self.optimizer = tf.keras.optimizers.Adam()
      self.train_loss = tf.keras.metrics.Mean(name="train_loss")
      self.train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="train_accuracy"
      )
  
      self.test_loss = tf.keras.metrics.Mean(name="test_loss")
      self.test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
          name="test_accuracy"
      )
  
      @tf.function
      def train_step(images, labels):
          with tf.GradientTape() as tape:
              predictions = self.model(images)
              loss = self.loss_object(labels, predictions)
          gradients = tape.gradient(
              loss, self.model.trainable_variables
          )
          self.optimizer.apply_gradients(
              zip(gradients, self.model.trainable_variables)
          )
  
          self.train_loss(loss)
          self.train_accuracy(labels, predictions)
  
      @tf.function
      def test_step(images, labels):
          predictions = self.model(images)
          t_loss = self.loss_object(labels, predictions)
  
          self.test_loss(t_loss)
          self.test_accuracy(labels, predictions)
  
      self.tf_train_step = train_step
      self.tf_test_step = test_step
  
file_path: ray/python/ray/tune/examples/tf_mnist_example.py
imports: import argparse
  import os
  from filelock import FileLock
  from tensorflow.keras.layers import Dense, Flatten, Conv2D
  from tensorflow.keras import Model
  from tensorflow.keras.datasets.mnist import load_data
  from ray import train, tune
label_desc: Use image node_id 1 for calling the MNISTTrainable obj's underlying member method code with example usage: obj.setup(config) without return types