title: TensorFlow Checkpoint restoring Learning rate
body: <p>I am trying to use <code>TensorFlow checkpoint</code>, everything is working well except for the <code>Learning rate</code>. It is getting re-initialized every time I run and does not restore from the previous.</p>
  <p>Here is a toy example I am trying to replicate the problem:</p>
  <pre><code>import numpy as np
  import tensorflow as tf
    
  
  X = tf.range(10.)
  Y = 50.*X
  
  class CGMM(tf.Module):
      def __init__(self):
          super(CGMM, self).__init__(name='CGMM')
          self.beta =  tf.Variable(1. , dtype=np.float32)
          self.learning_rate =  tf.Variable(1. , dtype=np.float32)
  
      @tf.function
      def objfun(self):
          beta = self.beta
          obj = tf.reduce_mean(tf.square(beta*self.X - self.Y))
          return obj
  
      def build_model(self,X,Y,decay_steps):
  
          self.X,self.Y=X,Y
          starter_learning_rate = 0.05 #0.05
          global_step = tf.Variable(1, trainable=False)
          self.learning_rate = tf.compat.v1.train.exponential_decay(starter_learning_rate, global_step,decay_steps, 0.96, staircase=True)
  
          optimizer = tf.compat.v1.train.RMSPropOptimizer(self.learning_rate)
  
          ckpt = tf.train.Checkpoint(step=tf.Variable(1) ,model=self, optimizer=optimizer)
          manager = tf.train.CheckpointManager(ckpt, './tf_ckpts_cg', max_to_keep=3)
  
          ckpt.restore(manager.latest_checkpoint)
          if manager.latest_checkpoint:
              print(&quot;Restored from {}&quot;.format(manager.latest_checkpoint))
          else:
              print(&quot;Initializing from scratch.&quot;)
  
          for i in range(500):
              optimizer.minimize(self.objfun, global_step=global_step,  var_list =  [self.beta])
              loss, beta, learning_rate = self.objfun(), self.beta, self.learning_rate().numpy()
  
              ckpt.step.assign_add(1)
              if (int(ckpt.step)-1) % 100 == 0:
                  save_path = manager.save()
                  print(&quot;Saved checkpoint for step {}: {}&quot;.format(int(ckpt.step), save_path))
                  print(&quot;learning_rate : &quot; + str(learning_rate))
  
          return beta
  
  
  model = CGMM()
  opt_beta = model.build_model(X,Y,100)
  </code></pre>
  <p>Results 1st run:</p>
  <pre><code>Initializing from scratch.
  Saved checkpoint for step 101: ./tf_ckpts_cg/ckpt-1
  learning_rate : 0.048
  Saved checkpoint for step 201: ./tf_ckpts_cg/ckpt-2
  learning_rate : 0.04608
  Saved checkpoint for step 301: ./tf_ckpts_cg/ckpt-3
  learning_rate : 0.044236798
  Saved checkpoint for step 401: ./tf_ckpts_cg/ckpt-4
  learning_rate : 0.042467322
  Saved checkpoint for step 501: ./tf_ckpts_cg/ckpt-5
  learning_rate : 0.04076863
  </code></pre>
  <p>Results 2nd run:</p>
  <pre><code>Restored from ./tf_ckpts_cg/ckpt-5
  Saved checkpoint for step 601: ./tf_ckpts_cg/ckpt-6
  learning_rate : 0.048
  Saved checkpoint for step 701: ./tf_ckpts_cg/ckpt-7
  learning_rate : 0.04608
  Saved checkpoint for step 801: ./tf_ckpts_cg/ckpt-8
  learning_rate : 0.044236798
  Saved checkpoint for step 901: ./tf_ckpts_cg/ckpt-9
  learning_rate : 0.042467322
  Saved checkpoint for step 1001: ./tf_ckpts_cg/ckpt-10
  learning_rate : 0.04076863
  </code></pre>
  <p>As you can see the same <code>Learning Rate</code> is repeated on both the runs but other variables are working well. Can you please help me to fix this?</p>