diff --git a/alibi_detect/saving/_tensorflow/loading.py b/alibi_detect/saving/_tensorflow/loading.py
index 1dc50b542..ab3410228 100644
--- a/alibi_detect/saving/_tensorflow/loading.py
+++ b/alibi_detect/saving/_tensorflow/loading.py
@@ -130,7 +130,7 @@ def load_kernel_config(cfg: dict) -> Callable:
     return kernel
 
 
-def load_optimizer(cfg: dict) -> Union[Type[tf.keras.optimizers.Optimizer], tf.keras.optimizers.Optimizer]:
+def load_optimizer(cfg: dict) -> Union[Type['tf.keras.optimizers.Optimizer'], 'tf.keras.optimizers.Optimizer']:
     """
     Loads a TensorFlow optimzier from a optimizer config dict.
 
diff --git a/alibi_detect/saving/_tensorflow/saving.py b/alibi_detect/saving/_tensorflow/saving.py
index cca7c96c6..9fb4e202d 100644
--- a/alibi_detect/saving/_tensorflow/saving.py
+++ b/alibi_detect/saving/_tensorflow/saving.py
@@ -158,7 +158,7 @@ def save_embedding_config(embed: TransformerEmbedding,
     return cfg_embed
 
 
-def save_optimizer_config(optimizer: Union[tf.keras.optimizers.Optimizer, tf.keras.optimizers.legacy.Optimizer]):
+def save_optimizer_config(optimizer: Union['tf.keras.optimizers.Optimizer', 'tf.keras.optimizers.legacy.Optimizer']):
     """
 
     Parameters