From 69b78e6623eeb9af2e47b0da2435ac840a9b810d Mon Sep 17 00:00:00 2001
From: bhack <bhack@users.noreply.github.com>
Date: Thu, 27 Aug 2020 16:14:27 +0200
Subject: [PATCH] Lookahead tracking base_optimizer

Inital fix of
https://github.com/tensorflow/addons/issues/2094
https://github.com/tensorflow/addons/pull/2102
---
 tensorflow_addons/optimizers/lookahead.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorflow_addons/optimizers/lookahead.py b/tensorflow_addons/optimizers/lookahead.py
index 8f96dc9c62..a1a4db9566 100644
--- 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(optimizer, 'base_optimizer')
 
     def _create_slots(self, var_list):
         self._optimizer._create_slots(