Failing to restore AdaBelief optimizer from checkpoint (#2705)

* Update adabelief.py

* addressed pr comments

Author: denadai2

tensorflow_addons/optimizers/adabelief.py

@@ -140,7 +140,7 @@ def __init__(
         self._set_hyper("decay", self._initial_decay)
         self._set_hyper("weight_decay", weight_decay)
         self._set_hyper("sma_threshold", sma_threshold)
-        self._set_hyper("total_steps", int(total_steps))
+        self._set_hyper("total_steps", float(total_steps))
         self._set_hyper("warmup_proportion", warmup_proportion)
         self._set_hyper("min_lr", min_lr)
         self.epsilon = epsilon or tf.keras.backend.epsilon()
@@ -325,7 +325,7 @@ def get_config(self):
                 "epsilon": self.epsilon,
                 "amsgrad": self.amsgrad,
                 "rectify": self.rectify,
-                "total_steps": self._serialize_hyperparameter("total_steps"),
+                "total_steps": int(self._serialize_hyperparameter("total_steps")),
                 "warmup_proportion": self._serialize_hyperparameter(
                     "warmup_proportion"
                 ),

tensorflow_addons/optimizers/tests/adabelief_test.py

@@ -236,3 +236,26 @@ def test_scheduler_serialization():
         "class_name": "InverseTimeDecay",
         "config": wd_scheduler.get_config(),
     }
+
+
+def test_checkpoint_serialization(tmpdir):
+    optimizer = AdaBelief()
+    optimizer2 = AdaBelief()
+
+    var_0 = tf.Variable([1.0, 2.0], dtype=tf.dtypes.float32)
+    var_1 = tf.Variable([3.0, 4.0], dtype=tf.dtypes.float32)
+
+    grad_0 = tf.constant([0.1, 0.2], dtype=tf.dtypes.float32)
+    grad_1 = tf.constant([0.03, 0.04], dtype=tf.dtypes.float32)
+
+    grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
+
+    optimizer.apply_gradients(grads_and_vars)
+
+    checkpoint = tf.train.Checkpoint(optimizer=optimizer)
+    checkpoint2 = tf.train.Checkpoint(optimizer=optimizer2)
+    model_path = str(tmpdir / "adabelief_chkpt")
+    checkpoint.write(model_path)
+    checkpoint2.read(model_path)
+
+    optimizer2.apply_gradients(grads_and_vars)