title: How should Exponential Moving Average be used in custom TF2.4 training loop
body: <p>I have a custom training loop that can be simplified as follow</p>
  <pre><code>inputs = tf.keras.Input(dtype=tf.float32, shape=(None, None, 3))
  model = tf.keras.Model({&quot;inputs&quot;: inputs}, {&quot;loss&quot;: f(inputs)})
  optimizer = tf.keras.optimizers.SGD(learning_rate=0.1, momentum=0.9, nesterov=True)
  
  for inputs in batches:
      with tf.GradientTape() as tape:
          results = model(inputs, training=True)
      grads = tape.gradient(results[&quot;loss&quot;], model.trainable_weights)
      optimizer.apply_gradients(zip(grads, model.trainable_weights))
  </code></pre>
  <p>The <a href="https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage" rel="nofollow noreferrer">TensorFlow documentation of ExponentialMovingAverage</a> is not clear on how it should be used in <a href="https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch" rel="nofollow noreferrer">from-scratch training loop</a>. As anyone worked with this?</p>
  <p>Additionally, how should the shadow variable be restored into the model if both are still in memory, and how can I check that that training variables were correctly updated?</p>