title: Steps of tf.summary.* operations in TensorBoard are always 0
body: <p>When I'm training my model with TensorFlow 2.3, I want to visualize some intermediate tensors calculated using the weight in the computation graph of my customized <code>tf.keras.layers.Layer</code>.</p>
  <p>So I use <code>tf.summary.image()</code> to record these tensors and visualize them as images like this:</p>
  <pre><code>class CustomizedLayer(tf.keras.layers.Layer):
      def call(self, inputs, training=None):
          # ... some code ...
          tf.summary.image(name=&quot;some_weight_map&quot;, data=some_weight_map)
          # ... some code ...
  </code></pre>
  <p>But in TensorBoard, no matter how many steps passed, there is only one image of step 0 shown.</p>
  <p>And I tried to set the parameter <em><strong>step</strong></em> of <code>tf.summary.image()</code> to the value obtained from <code>tf.summary.experimental.get_step()</code>:</p>
  <pre><code>tf.summary.image(name=&quot;weight_map&quot;, data=weight_map, step=tf.summary.experimental.get_step())
  </code></pre>
  <p>And update the step by calling <strong>tf.summary.experimental.set_step</strong> from a customized Callback using a tf.Variable like codes shown below:</p>
  <pre><code>class SummaryCallback(tf.keras.callbacks.Callback):
  def __init__(self, step_per_epoch):
      super().__init__()
      self.global_step = tf.Variable(initial_value=0, trainable=False, name=&quot;global_step&quot;)
      self.global_epoch = 0
      self.step_per_epoch = step_per_epoch
      tf.summary.experimental.set_step(self.global_step)
  
  def on_batch_end(self, batch, logs=None):
      self.global_step = batch + self.step_per_epoch * self.global_epoch
      tf.summary.experimental.set_step(self.global_step)  
      # whether the line above is commented, calling tf.summary.experimental.get_step() in computation graph code always returns 0.
      # tf.print(self.global_step)
  
  def on_epoch_end(self, epoch, logs=None):
      self.global_epoch += 1
  </code></pre>
  <p>This Callback's instance is passed in the argument <em><strong>callbacks</strong></em> in <code>model.fit()</code> function.</p>
  <p>But the value <code>tf.summary.experimental.get_step()</code> returned is still 0.</p>
  <p>The TensorFlow document of &quot;<code>tf.summary.experimental.set_step()</code>&quot; says:</p>
  <blockquote>
  <p>when using this with @tf.functions, the step value will be captured at the time the function is traced, so changes to the step outside the function will not be reflected inside the function unless using a tf.Variable step.</p>
  </blockquote>
  <p>Accroding to the document, I am already using a Variable to store the steps, but it's changes are still not reflected inside the function (or keras.Model).</p>
  <p>Note: My code produces expected results in TensorFlow 1.x with just a simple line of <code>tf.summary.image()</code> before I migrate it to TensorFlow 2.</p>
  <p>So I want to know if my approach is wrong in TensorFlow 2?</p>
  <p>In TF2, how can I <strong>get training steps inside the computation graph</strong>?</p>
  <p>Or there is other solution to <strong>summarize tensors (as scalar, image, etc.) inside a model in TensorFlow 2</strong>?</p>