title: Tensorboard Graph with custom training loop does not include my Model
body: <p>I have created my own loop as shown in the TF 2 migration guide <a href="https://www.tensorflow.org/guide/migrate#write_your_own_loop" rel="nofollow noreferrer">here</a>.<br>
  I am currently able to see the graph for only the <code>--- VISIBLE ---</code> section of the code below. How do I make my model (defined in the <code>---NOT VISIBLE---</code> section) visible in tensorboard?</p>
  
  <p>If I was not using a custom training loop, I could have gone with the <a href="https://www.tensorflow.org/tensorboard/graphs#train_the_model_and_log_data" rel="nofollow noreferrer">documented</a> <code>model.fit approach</code>:</p>
  
  <pre><code>model.fit(..., callbacks=[keras.callbacks.TensorBoard(log_dir=logdir)])
  </code></pre>
  
  <p>In TF 1, the approach used to be quite straightforward:</p>
  
  <pre><code>tf.compat.v1.summary.FileWriter(LOGDIR, sess.graph)
  </code></pre>
  
  <p>The Tensorboard migration guide clearly states (<a href="https://www.tensorflow.org/tensorboard/migrate#additional_tips" rel="nofollow noreferrer">here</a>) that:</p>
  
  <blockquote>
    <p>No direct writing of tf.compat.v1.Graph - instead use @tf.function and trace functions</p>
  </blockquote>
  
  <pre><code>configure_default_gpus()
  tf.summary.trace_on(graph=True)
  K = tf.keras
  dataset = sanity_dataset(BATCH_SIZE)
  
  #-------------------------- NOT VISIBLE -----------------------------------------
  model = K.models.Sequential([
      K.layers.Flatten(input_shape=(IMG_WIDTH, IMG_HEIGHT, IMG_CHANNELS)),
      K.layers.Dense(10, activation=K.layers.LeakyReLU()),
      K.layers.Dense(IMG_WIDTH * IMG_HEIGHT * IMG_CHANNELS, activation=K.layers.LeakyReLU()),
      K.layers.Reshape((IMG_WIDTH, IMG_HEIGHT, IMG_CHANNELS)),
  ])
  #--------------------------------------------------------------------------------
  
  optimizer = tf.keras.optimizers.Adam()
  loss_fn = K.losses.Huber()
  
  
  @tf.function
  def train_step(inputs, targets):
      with tf.GradientTape() as tape:
          predictions = model(inputs, training=True)
  #-------------------------- VISIBLE ---------------------------------------------
          pred_loss = loss_fn(targets, predictions)
  
      gradients = tape.gradient(pred_loss, model.trainable_variables)
      optimizer.apply_gradients(zip(gradients, model.trainable_variables))
  #--------------------------------------------------------------------------------
      return pred_loss, predictions
  
  
  with tf.summary.create_file_writer(LOG_DIR).as_default() as writer:
      for epoch in range(5):
          for step, (input_batch, target_batch) in enumerate(dataset):
              total_loss, predictions = train_step(input_batch, target_batch)
  
              if step == 0:
                  tf.summary.trace_export(name="all", step=step, profiler_outdir=LOG_DIR)
              tf.summary.scalar('loss', total_loss, step=step)
              writer.flush()
  writer.close()
  </code></pre>
  
  <p>There's a <a href="https://stackoverflow.com/questions/60639731/tensorboard-for-custom-training-loop-in-tensorflow-2">similar unanswered question</a> where the OP was unable to view any graph.</p>