Update lookahead.py

Inital fix of #2094 #2102
tensorflow · Aug 27, 2020 · 0a57c39 · 0a57c39
1 parent 6329e1a
commit 0a57c39
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py
@@ -80,6 +80,7 @@ def __init__(
         self._set_hyper("sync_period", sync_period)
         self._set_hyper("slow_step_size", slow_step_size)
         self._initialized = False
+        self._track_trackable(self._optimizer, 'base_optimizer')
 
     def _create_slots(self, var_list):
         self._optimizer._create_slots(