From 169942255c3604e6b0bfbb53a3313652a29c8974 Mon Sep 17 00:00:00 2001
From: veraAI <113360908+veraAI@users.noreply.github.com>
Date: Mon, 12 Sep 2022 16:08:29 +0200
Subject: [PATCH] Fixed TypeError in CyclicLearningRate

Fixed bug which appears at the end of the first epoch. With TF2.5 or higher the optimizer returns the step variable as int64 but the scale_fn function needs it to be float for Triangular2CyclicalLearningRate and ExponentialCyclicalLearningRate.
---
 tensorflow_addons/optimizers/cyclical_learning_rate.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow_addons/optimizers/cyclical_learning_rate.py b/tensorflow_addons/optimizers/cyclical_learning_rate.py
index f4ec176977..be1ca59df8 100644
--- a/tensorflow_addons/optimizers/cyclical_learning_rate.py
+++ b/tensorflow_addons/optimizers/cyclical_learning_rate.py
@@ -95,7 +95,7 @@ def __call__(self, step):
             cycle = tf.floor(1 + step_as_dtype / (2 * step_size))
             x = tf.abs(step_as_dtype / step_size - 2 * cycle + 1)
 
-            mode_step = cycle if self.scale_mode == "cycle" else step
+            mode_step = cycle if self.scale_mode == "cycle" else step_as_dtype
 
             return initial_learning_rate + (
                 maximal_learning_rate - initial_learning_rate