diff --git a/docs/source/using_doctr/using_model_export.rst b/docs/source/using_doctr/using_model_export.rst
index 48f570f699..c62c36169b 100644
--- a/docs/source/using_doctr/using_model_export.rst
+++ b/docs/source/using_doctr/using_model_export.rst
@@ -31,7 +31,7 @@ Advantages:
         .. code:: python3
 
             import tensorflow as tf
-            from keras import mixed_precision
+            from tensorflow.keras import mixed_precision
             mixed_precision.set_global_policy('mixed_float16')
             predictor = ocr_predictor(reco_arch="crnn_mobilenet_v3_small", det_arch="linknet_resnet34", pretrained=True)
 
diff --git a/doctr/file_utils.py b/doctr/file_utils.py
index 68e9dfffac..fc1129b0c1 100644
--- a/doctr/file_utils.py
+++ b/doctr/file_utils.py
@@ -35,6 +35,20 @@
     logging.info("Disabling PyTorch because USE_TF is set")
     _torch_available = False
 
+# Compatibility fix to make sure tensorflow.keras stays at Keras 2
+if "TF_USE_LEGACY_KERAS" not in os.environ:
+    os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
+elif os.environ["TF_USE_LEGACY_KERAS"] != "1":
+    raise ValueError(
+        "docTR is only compatible with Keras 2, but you have explicitly set `TF_USE_LEGACY_KERAS` to `0`. "
+    )
+
+
+def ensure_keras_v2() -> None:  # pragma: no cover
+    if not os.environ.get("TF_USE_LEGACY_KERAS") == "1":
+        os.environ["TF_USE_LEGACY_KERAS"] = "1"
+
 
 if USE_TF in ENV_VARS_TRUE_AND_AUTO_VALUES and USE_TORCH not in ENV_VARS_TRUE_VALUES:
     _tf_available = importlib.util.find_spec("tensorflow") is not None
@@ -65,6 +79,11 @@
             _tf_available = False
         else:
             logging.info(f"TensorFlow version {_tf_version} available.")
+            ensure_keras_v2()
+            import tensorflow as tf
+
+            # Enable eager execution - this is required for some models to work properly
+            tf.config.run_functions_eagerly(True)
 else:  # pragma: no cover
     logging.info("Disabling Tensorflow because USE_TORCH is set")
     _tf_available = False
diff --git a/doctr/io/image/tensorflow.py b/doctr/io/image/tensorflow.py
index 3b1f1ed0e2..28fb2fadd5 100644
--- a/doctr/io/image/tensorflow.py
+++ b/doctr/io/image/tensorflow.py
@@ -7,8 +7,8 @@
 
 import numpy as np
 import tensorflow as tf
-from keras.utils import img_to_array
 from PIL import Image
+from tensorflow.keras.utils import img_to_array
 
 from doctr.utils.common_types import AbstractPath
 
diff --git a/doctr/models/classification/magc_resnet/tensorflow.py b/doctr/models/classification/magc_resnet/tensorflow.py
index 12f7c6beea..fc7678f661 100644
--- a/doctr/models/classification/magc_resnet/tensorflow.py
+++ b/doctr/models/classification/magc_resnet/tensorflow.py
@@ -9,8 +9,8 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import activations, layers
-from keras.models import Sequential
+from tensorflow.keras import activations, layers
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/classification/mobilenet/tensorflow.py b/doctr/models/classification/mobilenet/tensorflow.py
index 6250abc666..ff57c221dc 100644
--- a/doctr/models/classification/mobilenet/tensorflow.py
+++ b/doctr/models/classification/mobilenet/tensorflow.py
@@ -9,8 +9,8 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import tensorflow as tf
-from keras import layers
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Sequential
 
 from ....datasets import VOCABS
 from ...utils import conv_sequence, load_pretrained_params
diff --git a/doctr/models/classification/predictor/tensorflow.py b/doctr/models/classification/predictor/tensorflow.py
index ba26e1db54..23efbf6579 100644
--- a/doctr/models/classification/predictor/tensorflow.py
+++ b/doctr/models/classification/predictor/tensorflow.py
@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model
+from tensorflow.keras import Model
 
 from doctr.models.preprocessor import PreProcessor
 from doctr.utils.repr import NestedObject
diff --git a/doctr/models/classification/resnet/tensorflow.py b/doctr/models/classification/resnet/tensorflow.py
index 3e78ae0ae2..364b03c3a2 100644
--- a/doctr/models/classification/resnet/tensorflow.py
+++ b/doctr/models/classification/resnet/tensorflow.py
@@ -7,9 +7,9 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import layers
-from keras.applications import ResNet50
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.applications import ResNet50
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/classification/textnet/tensorflow.py b/doctr/models/classification/textnet/tensorflow.py
index 3d79b15f09..b0bb9a7205 100644
--- a/doctr/models/classification/textnet/tensorflow.py
+++ b/doctr/models/classification/textnet/tensorflow.py
@@ -7,7 +7,7 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple
 
-from keras import Sequential, layers
+from tensorflow.keras import Sequential, layers
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/classification/vgg/tensorflow.py b/doctr/models/classification/vgg/tensorflow.py
index d9e7bb374b..9ecdabd040 100644
--- a/doctr/models/classification/vgg/tensorflow.py
+++ b/doctr/models/classification/vgg/tensorflow.py
@@ -6,8 +6,8 @@
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Tuple
 
-from keras import layers
-from keras.models import Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Sequential
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/classification/vit/tensorflow.py b/doctr/models/classification/vit/tensorflow.py
index 28ff2e244e..8531193939 100644
--- a/doctr/models/classification/vit/tensorflow.py
+++ b/doctr/models/classification/vit/tensorflow.py
@@ -7,7 +7,7 @@
 from typing import Any, Dict, Optional, Tuple
 
 import tensorflow as tf
-from keras import Sequential, layers
+from tensorflow.keras import Sequential, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.modules.transformer import EncoderBlock
diff --git a/doctr/models/detection/differentiable_binarization/tensorflow.py b/doctr/models/detection/differentiable_binarization/tensorflow.py
index 7fdbd43ce0..45e522b872 100644
--- a/doctr/models/detection/differentiable_binarization/tensorflow.py
+++ b/doctr/models/detection/differentiable_binarization/tensorflow.py
@@ -10,8 +10,8 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers, losses
-from keras.applications import ResNet50
+from tensorflow.keras import Model, Sequential, layers, losses
+from tensorflow.keras.applications import ResNet50
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, conv_sequence, load_pretrained_params
diff --git a/doctr/models/detection/fast/tensorflow.py b/doctr/models/detection/fast/tensorflow.py
index 80fc31fea3..91d6c8cc4d 100644
--- a/doctr/models/detection/fast/tensorflow.py
+++ b/doctr/models/detection/fast/tensorflow.py
@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers
+from tensorflow.keras import Model, Sequential, layers
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.utils import IntermediateLayerGetter, _bf16_to_float32, load_pretrained_params
diff --git a/doctr/models/detection/linknet/tensorflow.py b/doctr/models/detection/linknet/tensorflow.py
index 683c49373a..df8233cf20 100644
--- a/doctr/models/detection/linknet/tensorflow.py
+++ b/doctr/models/detection/linknet/tensorflow.py
@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, Sequential, layers, losses
+from tensorflow.keras import Model, Sequential, layers, losses
 
 from doctr.file_utils import CLASS_NAME
 from doctr.models.classification import resnet18, resnet34, resnet50
diff --git a/doctr/models/detection/predictor/tensorflow.py b/doctr/models/detection/predictor/tensorflow.py
index a7ccd4a9ac..a3d5085847 100644
--- a/doctr/models/detection/predictor/tensorflow.py
+++ b/doctr/models/detection/predictor/tensorflow.py
@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model
+from tensorflow.keras import Model
 
 from doctr.models.detection._utils import _remove_padding
 from doctr.models.preprocessor import PreProcessor
diff --git a/doctr/models/modules/layers/tensorflow.py b/doctr/models/modules/layers/tensorflow.py
index b1019be778..68849fbf6e 100644
--- a/doctr/models/modules/layers/tensorflow.py
+++ b/doctr/models/modules/layers/tensorflow.py
@@ -7,7 +7,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 
diff --git a/doctr/models/modules/transformer/tensorflow.py b/doctr/models/modules/transformer/tensorflow.py
index eef4f3dbea..50c7cef04d 100644
--- a/doctr/models/modules/transformer/tensorflow.py
+++ b/doctr/models/modules/transformer/tensorflow.py
@@ -7,14 +7,12 @@
 from typing import Any, Callable, Optional, Tuple
 
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 
 __all__ = ["Decoder", "PositionalEncoding", "EncoderBlock", "PositionwiseFeedForward", "MultiHeadAttention"]
 
-tf.config.run_functions_eagerly(True)
-
 
 class PositionalEncoding(layers.Layer, NestedObject):
     """Compute positional encoding"""
diff --git a/doctr/models/modules/vision_transformer/tensorflow.py b/doctr/models/modules/vision_transformer/tensorflow.py
index a73aa4c706..8386172eb1 100644
--- a/doctr/models/modules/vision_transformer/tensorflow.py
+++ b/doctr/models/modules/vision_transformer/tensorflow.py
@@ -7,7 +7,7 @@
 from typing import Any, Tuple
 
 import tensorflow as tf
-from keras import layers
+from tensorflow.keras import layers
 
 from doctr.utils.repr import NestedObject
 
diff --git a/doctr/models/recognition/crnn/tensorflow.py b/doctr/models/recognition/crnn/tensorflow.py
index d366bfc14b..fb5cb72dff 100644
--- a/doctr/models/recognition/crnn/tensorflow.py
+++ b/doctr/models/recognition/crnn/tensorflow.py
@@ -7,8 +7,8 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import tensorflow as tf
-from keras import layers
-from keras.models import Model, Sequential
+from tensorflow.keras import layers
+from tensorflow.keras.models import Model, Sequential
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/recognition/master/tensorflow.py b/doctr/models/recognition/master/tensorflow.py
index 5b8192dee6..42cd216b2c 100644
--- a/doctr/models/recognition/master/tensorflow.py
+++ b/doctr/models/recognition/master/tensorflow.py
@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.classification import magc_resnet31
diff --git a/doctr/models/recognition/parseq/tensorflow.py b/doctr/models/recognition/parseq/tensorflow.py
index bca7806903..b0e21a50d6 100644
--- a/doctr/models/recognition/parseq/tensorflow.py
+++ b/doctr/models/recognition/parseq/tensorflow.py
@@ -10,7 +10,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 from doctr.models.modules.transformer import MultiHeadAttention, PositionwiseFeedForward
@@ -167,7 +167,6 @@ def __init__(
 
         self.postprocessor = PARSeqPostProcessor(vocab=self.vocab)
 
-    @tf.function
     def generate_permutations(self, seqlen: tf.Tensor) -> tf.Tensor:
         # Generates permutations of the target sequence.
         # Translated from https://github.com/baudm/parseq/blob/main/strhub/models/parseq/system.py
@@ -214,7 +213,6 @@ def generate_permutations(self, seqlen: tf.Tensor) -> tf.Tensor:
             )
         return combined
 
-    @tf.function
     def generate_permutations_attention_masks(self, permutation: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor]:
         # Generate source and target mask for the decoder attention.
         sz = permutation.shape[0]
@@ -234,7 +232,6 @@ def generate_permutations_attention_masks(self, permutation: tf.Tensor) -> Tuple
         target_mask = mask[1:, :-1]
         return tf.cast(source_mask, dtype=tf.bool), tf.cast(target_mask, dtype=tf.bool)
 
-    @tf.function
     def decode(
         self,
         target: tf.Tensor,
diff --git a/doctr/models/recognition/sar/tensorflow.py b/doctr/models/recognition/sar/tensorflow.py
index 0776414c7a..89e93ea51e 100644
--- a/doctr/models/recognition/sar/tensorflow.py
+++ b/doctr/models/recognition/sar/tensorflow.py
@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, Sequential, layers
+from tensorflow.keras import Model, Sequential, layers
 
 from doctr.datasets import VOCABS
 from doctr.utils.repr import NestedObject
diff --git a/doctr/models/recognition/vitstr/tensorflow.py b/doctr/models/recognition/vitstr/tensorflow.py
index 985f49a470..6b38cf7548 100644
--- a/doctr/models/recognition/vitstr/tensorflow.py
+++ b/doctr/models/recognition/vitstr/tensorflow.py
@@ -7,7 +7,7 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 import tensorflow as tf
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.datasets import VOCABS
 
diff --git a/doctr/models/utils/tensorflow.py b/doctr/models/utils/tensorflow.py
index 51a2bc69a5..6f7dc14ab3 100644
--- a/doctr/models/utils/tensorflow.py
+++ b/doctr/models/utils/tensorflow.py
@@ -8,7 +8,7 @@
 
 import tensorflow as tf
 import tf2onnx
-from keras import Model, layers
+from tensorflow.keras import Model, layers
 
 from doctr.utils.data import download_from_url
 
@@ -77,7 +77,7 @@ def conv_sequence(
 ) -> List[layers.Layer]:
     """Builds a convolutional-based layer sequence
 
-    >>> from keras import Sequential
+    >>> from tensorflow.keras import Sequential
     >>> from doctr.models import conv_sequence
     >>> module = Sequential(conv_sequence(32, 'relu', True, kernel_size=3, input_shape=[224, 224, 3]))
 
@@ -113,7 +113,7 @@ def conv_sequence(
 class IntermediateLayerGetter(Model):
     """Implements an intermediate layer getter
 
-    >>> from keras.applications import ResNet50
+    >>> from tensorflow.keras.applications import ResNet50
     >>> from doctr.models import IntermediateLayerGetter
     >>> target_layers = ["conv2_block3_out", "conv3_block4_out", "conv4_block6_out", "conv5_block3_out"]
     >>> feat_extractor = IntermediateLayerGetter(ResNet50(include_top=False, pooling=False), target_layers)
diff --git a/doctr/transforms/modules/tensorflow.py b/doctr/transforms/modules/tensorflow.py
index b3f7bcfd8a..8dd870e3e7 100644
--- a/doctr/transforms/modules/tensorflow.py
+++ b/doctr/transforms/modules/tensorflow.py
@@ -395,7 +395,6 @@ def __init__(self, kernel_shape: Union[int, Iterable[int]], std: Tuple[float, fl
     def extra_repr(self) -> str:
         return f"kernel_shape={self.kernel_shape}, std={self.std}"
 
-    @tf.function
     def __call__(self, img: tf.Tensor) -> tf.Tensor:
         return tf.squeeze(
             _gaussian_filter(
diff --git a/pyproject.toml b/pyproject.toml
index aa0e02f98e..9745f8a7c4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -52,12 +52,10 @@ dependencies = [
 
 [project.optional-dependencies]
 tf = [
-    # cf. https://github.com/mindee/doctr/pull/1182
     # cf. https://github.com/mindee/doctr/pull/1461
-    "tensorflow>=2.11.0,<2.16.0",
+    "tensorflow>=2.15.0,<3.0.0",
+    "tf-keras>=2.15.0,<3.0.0",  # Keep keras 2 compatibility
     "tf2onnx>=1.16.0,<2.0.0",  # cf. https://github.com/onnx/tensorflow-onnx/releases/tag/v1.16.0
-    # TODO: This is a temporary fix until we can upgrade to a newer version of tensorflow
-    "numpy>=1.16.0,<2.0.0",
 ]
 torch = [
     "torch>=1.12.0,<3.0.0",
@@ -98,9 +96,9 @@ docs = [
 ]
 dev = [
     # Tensorflow
-    # cf. https://github.com/mindee/doctr/pull/1182
     # cf. https://github.com/mindee/doctr/pull/1461
-    "tensorflow>=2.11.0,<2.16.0",
+    "tensorflow>=2.15.0,<3.0.0",
+    "tf-keras>=2.15.0,<3.0.0",  # Keep keras 2 compatibility
     "tf2onnx>=1.16.0,<2.0.0",  # cf. https://github.com/onnx/tensorflow-onnx/releases/tag/v1.16.0
     # PyTorch
     "torch>=1.12.0,<3.0.0",
diff --git a/references/classification/README.md b/references/classification/README.md
index d0e5c3b83a..6646b0d8ca 100644
--- a/references/classification/README.md
+++ b/references/classification/README.md
@@ -60,12 +60,12 @@ Feel free to inspect the multiple script option to customize your training to yo
 
 Character classification:
 
-```python
+```shell
 python references/classification/train_tensorflow_character.py --help
 ```
 
 Orientation classification:
 
-```python
+```shell
 python references/classification/train_tensorflow_orientation.py --help
 ```
diff --git a/references/classification/latency_tensorflow.py b/references/classification/latency_tensorflow.py
index fc010df91a..5a1d3f7845 100644
--- a/references/classification/latency_tensorflow.py
+++ b/references/classification/latency_tensorflow.py
@@ -9,12 +9,16 @@
 import os
 import time
 
-import numpy as np
-import tensorflow as tf
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
 
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
+import numpy as np
+import tensorflow as tf
+
 from doctr.models import classification
 
 
diff --git a/references/classification/train_tensorflow_character.py b/references/classification/train_tensorflow_character.py
index b2d24f2dbf..2c754df3d4 100644
--- a/references/classification/train_tensorflow_character.py
+++ b/references/classification/train_tensorflow_character.py
@@ -5,6 +5,10 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
@@ -13,7 +17,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub
@@ -82,6 +86,11 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
+@tf.function
+def apply_grads(optimizer, grads, model):
+    optimizer.apply_gradients(zip(grads, model.trainable_weights))
+
+
 def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     # Iterate over the batches of the dataset
     pbar = tqdm(train_loader, position=1)
@@ -94,7 +103,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
         grads = tape.gradient(train_loss, model.trainable_weights)
         if amp:
             grads = optimizer.get_unscaled_gradients(grads)
-        optimizer.apply_gradients(zip(grads, model.trainable_weights))
+        apply_grads(optimizer, grads, model)
 
         pbar.set_description(f"Training loss: {train_loss.numpy().mean():.6}")
 
diff --git a/references/classification/train_tensorflow_orientation.py b/references/classification/train_tensorflow_orientation.py
index e063174944..e29001e7bc 100644
--- a/references/classification/train_tensorflow_orientation.py
+++ b/references/classification/train_tensorflow_orientation.py
@@ -5,6 +5,10 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
@@ -13,7 +17,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub
@@ -96,6 +100,11 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
+@tf.function
+def apply_grads(optimizer, grads, model):
+    optimizer.apply_gradients(zip(grads, model.trainable_weights))
+
+
 def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     # Iterate over the batches of the dataset
     pbar = tqdm(train_loader, position=1)
@@ -108,7 +117,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
         grads = tape.gradient(train_loss, model.trainable_weights)
         if amp:
             grads = optimizer.get_unscaled_gradients(grads)
-        optimizer.apply_gradients(zip(grads, model.trainable_weights))
+        apply_grads(optimizer, grads, model)
 
         pbar.set_description(f"Training loss: {train_loss.numpy().mean():.6}")
 
diff --git a/references/detection/evaluate_tensorflow.py b/references/detection/evaluate_tensorflow.py
index abf012ed83..ea77037804 100644
--- a/references/detection/evaluate_tensorflow.py
+++ b/references/detection/evaluate_tensorflow.py
@@ -5,6 +5,10 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 from doctr.file_utils import CLASS_NAME
 
 os.environ["USE_TF"] = "1"
@@ -14,7 +18,7 @@
 from pathlib import Path
 
 import tensorflow as tf
-from keras import mixed_precision
+from tensorflow.keras import mixed_precision
 from tqdm import tqdm
 
 gpu_devices = tf.config.experimental.list_physical_devices("GPU")
diff --git a/references/detection/latency_tensorflow.py b/references/detection/latency_tensorflow.py
index e3e0d1d8af..b2d973fa6f 100644
--- a/references/detection/latency_tensorflow.py
+++ b/references/detection/latency_tensorflow.py
@@ -9,12 +9,17 @@
 import os
 import time
 
-import numpy as np
-import tensorflow as tf
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
 
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
+
+import numpy as np
+import tensorflow as tf
+
 from doctr.models import detection
 
 
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index b9c14494ad..a5a1caf2c5 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -5,6 +5,10 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
@@ -14,7 +18,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub
@@ -82,6 +86,11 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
+@tf.function
+def apply_grads(optimizer, grads, model):
+    optimizer.apply_gradients(zip(grads, model.trainable_weights))
+
+
 def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
@@ -94,7 +103,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
         grads = tape.gradient(train_loss, model.trainable_weights)
         if amp:
             grads = optimizer.get_unscaled_gradients(grads)
-        optimizer.apply_gradients(zip(grads, model.trainable_weights))
+        apply_grads(optimizer, grads, model)
 
         pbar.set_description(f"Training loss: {train_loss.numpy():.6}")
 
diff --git a/references/recognition/README.md b/references/recognition/README.md
index b82a0d99b5..5823030120 100644
--- a/references/recognition/README.md
+++ b/references/recognition/README.md
@@ -81,7 +81,7 @@ When typing your labels, be aware that the VOCAB doesn't handle spaces. Also mak
 
 Feel free to inspect the multiple script option to customize your training to your own needs!
 
-```python
+```shell
 python references/recognition/train_pytorch.py --help
 ```
 
diff --git a/references/recognition/evaluate_tensorflow.py b/references/recognition/evaluate_tensorflow.py
index 4c9d125285..4e00715513 100644
--- a/references/recognition/evaluate_tensorflow.py
+++ b/references/recognition/evaluate_tensorflow.py
@@ -5,13 +5,17 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
 import time
 
 import tensorflow as tf
-from keras import mixed_precision
+from tensorflow.keras import mixed_precision
 from tqdm import tqdm
 
 gpu_devices = tf.config.experimental.list_physical_devices("GPU")
diff --git a/references/recognition/latency_tensorflow.py b/references/recognition/latency_tensorflow.py
index 405cf56892..dfae409473 100644
--- a/references/recognition/latency_tensorflow.py
+++ b/references/recognition/latency_tensorflow.py
@@ -9,12 +9,16 @@
 import os
 import time
 
-import numpy as np
-import tensorflow as tf
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
 
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
+import numpy as np
+import tensorflow as tf
+
 from doctr.models import recognition
 
 
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
index c76355a2f2..ba560c285e 100644
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -5,6 +5,10 @@
 
 import os
 
+from doctr.file_utils import ensure_keras_v2
+
+ensure_keras_v2()
+
 os.environ["USE_TF"] = "1"
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
@@ -15,7 +19,7 @@
 
 import numpy as np
 import tensorflow as tf
-from keras import Model, mixed_precision, optimizers
+from tensorflow.keras import Model, mixed_precision, optimizers
 from tqdm.auto import tqdm
 
 from doctr.models import login_to_hub, push_to_hf_hub
@@ -83,6 +87,11 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
+@tf.function
+def apply_grads(optimizer, grads, model):
+    optimizer.apply_gradients(zip(grads, model.trainable_weights))
+
+
 def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
@@ -95,7 +104,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
         grads = tape.gradient(train_loss, model.trainable_weights)
         if amp:
             grads = optimizer.get_unscaled_gradients(grads)
-        optimizer.apply_gradients(zip(grads, model.trainable_weights))
+        apply_grads(optimizer, grads, model)
 
         pbar.set_description(f"Training loss: {train_loss.numpy().mean():.6}")
 
@@ -254,7 +263,7 @@ def main(args):
                 T.RandomSaturation(0.3),
                 T.RandomContrast(0.3),
                 T.RandomBrightness(0.3),
-                T.RandomApply(T.RandomShadow(), 0.4),
+                # T.RandomApply(T.RandomShadow(), 0.4),  # NOTE: RandomShadow is broken atm
                 T.RandomApply(T.GaussianNoise(mean=0.1, std=0.1), 0.1),
                 T.RandomApply(T.GaussianBlur(kernel_shape=3, std=(0.1, 0.1)), 0.3),
             ]),
diff --git a/tests/tensorflow/test_models_classification_tf.py b/tests/tensorflow/test_models_classification_tf.py
index 11f4ea4114..89181aace0 100644
--- a/tests/tensorflow/test_models_classification_tf.py
+++ b/tests/tensorflow/test_models_classification_tf.py
@@ -2,7 +2,6 @@
 import tempfile
 
 import cv2
-import keras
 import numpy as np
 import onnxruntime
 import psutil
@@ -38,7 +37,7 @@
 def test_classification_architectures(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     model = classification.__dict__[arch_name](pretrained=True, include_top=True, input_shape=input_shape)
     # Forward
     out = model(tf.random.uniform(shape=[batch_size, *input_shape], maxval=1, dtype=tf.float32))
@@ -47,7 +46,7 @@ def test_classification_architectures(arch_name, input_shape, output_size):
     assert out.dtype == tf.float32
     assert out.numpy().shape == (batch_size, *output_size)
     # Check that you can load pretrained up to the classification layer with differing number of classes to fine-tune
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     assert classification.__dict__[arch_name](
         pretrained=True, include_top=True, input_shape=input_shape, num_classes=10
     )
@@ -63,7 +62,7 @@ def test_classification_architectures(arch_name, input_shape, output_size):
 def test_classification_models(arch_name, input_shape):
     batch_size = 8
     reco_model = classification.__dict__[arch_name](pretrained=True, input_shape=input_shape)
-    assert isinstance(reco_model, keras.Model)
+    assert isinstance(reco_model, tf.keras.Model)
     input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
 
     out = reco_model(input_tensor)
@@ -232,7 +231,7 @@ def test_page_orientation_model(mock_payslip):
 def test_models_onnx_export(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if "orientation" in arch_name:
         model = classification.__dict__[arch_name](pretrained=True, input_shape=input_shape)
     else:
@@ -252,7 +251,6 @@ def test_models_onnx_export(arch_name, input_shape, output_size):
         model_path, output = export_model_to_onnx(
             model, model_name=os.path.join(tmpdir, "model"), dummy_input=dummy_input
         )
-
         assert os.path.exists(model_path)
         # Inference
         ort_session = onnxruntime.InferenceSession(
diff --git a/tests/tensorflow/test_models_detection_tf.py b/tests/tensorflow/test_models_detection_tf.py
index ba5f50542b..7dbb090bf2 100644
--- a/tests/tensorflow/test_models_detection_tf.py
+++ b/tests/tensorflow/test_models_detection_tf.py
@@ -2,7 +2,6 @@
 import os
 import tempfile
 
-import keras
 import numpy as np
 import onnxruntime
 import psutil
@@ -38,13 +37,13 @@
 )
 def test_detection_models(arch_name, input_shape, output_size, out_prob, train_mode):
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if arch_name == "fast_tiny_rep":
         model = reparameterize(detection.fast_tiny(pretrained=True, input_shape=input_shape))
         train_mode = False  # Reparameterized model is not trainable
     else:
         model = detection.__dict__[arch_name](pretrained=True, input_shape=input_shape)
-    assert isinstance(model, keras.Model)
+    assert isinstance(model, tf.keras.Model)
     input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
     target = [
         {CLASS_NAME: np.array([[0.5, 0.5, 1, 1], [0.5, 0.5, 0.8, 0.8]], dtype=np.float32)},
@@ -153,7 +152,7 @@ def test_rotated_detectionpredictor(mock_pdf):
 )
 def test_detection_zoo(arch_name):
     # Model
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     predictor = detection.zoo.detection_predictor(arch_name, pretrained=False)
     # object check
     assert isinstance(predictor, DetectionPredictor)
@@ -178,7 +177,7 @@ def test_fast_reparameterization():
     base_model_params = np.sum([np.prod(v.shape) for v in base_model.trainable_variables])
     assert math.isclose(base_model_params, 13535296)  # base model params
     base_out = base_model(dummy_input, training=False)["logits"]
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     rep_model = reparameterize(base_model)
     rep_model_params = np.sum([np.prod(v.shape) for v in base_model.trainable_variables])
     assert math.isclose(rep_model_params, 8520256)  # reparameterized model params
@@ -242,7 +241,7 @@ def test_dilate():
 def test_models_onnx_export(arch_name, input_shape, output_size):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     if arch_name == "fast_tiny_rep":
         model = reparameterize(detection.fast_tiny(pretrained=True, exportable=True, input_shape=input_shape))
     else:
@@ -257,6 +256,7 @@ def test_models_onnx_export(arch_name, input_shape, output_size):
             model, model_name=os.path.join(tmpdir, "model"), dummy_input=dummy_input
         )
         assert os.path.exists(model_path)
+
         # Inference
         ort_session = onnxruntime.InferenceSession(
             os.path.join(tmpdir, "model.onnx"), providers=["CPUExecutionProvider"]
diff --git a/tests/tensorflow/test_models_factory.py b/tests/tensorflow/test_models_factory.py
index 0860d8612c..a4483800c9 100644
--- a/tests/tensorflow/test_models_factory.py
+++ b/tests/tensorflow/test_models_factory.py
@@ -2,8 +2,8 @@
 import os
 import tempfile
 
-import keras
 import pytest
+import tensorflow as tf
 
 from doctr import models
 from doctr.models.factory import _save_model_and_config_for_hf_hub, from_hub, push_to_hf_hub
@@ -50,7 +50,7 @@ def test_push_to_hf_hub():
 )
 def test_models_for_hub(arch_name, task_name, dummy_model_id, tmpdir):
     with tempfile.TemporaryDirectory() as tmp_dir:
-        keras.backend.clear_session()
+        tf.keras.backend.clear_session()
         model = models.__dict__[task_name].__dict__[arch_name](pretrained=True)
 
         _save_model_and_config_for_hf_hub(model, arch=arch_name, task=task_name, save_dir=tmp_dir)
@@ -65,6 +65,6 @@ def test_models_for_hub(arch_name, task_name, dummy_model_id, tmpdir):
         assert all(key in model.cfg.keys() for key in tmp_config.keys())
 
         # test from hub
-        keras.backend.clear_session()
+        tf.keras.backend.clear_session()
         hub_model = from_hub(repo_id=dummy_model_id)
         assert isinstance(hub_model, type(model))
diff --git a/tests/tensorflow/test_models_recognition_tf.py b/tests/tensorflow/test_models_recognition_tf.py
index 7da1cb534a..162c446d35 100644
--- a/tests/tensorflow/test_models_recognition_tf.py
+++ b/tests/tensorflow/test_models_recognition_tf.py
@@ -2,7 +2,6 @@
 import shutil
 import tempfile
 
-import keras
 import numpy as np
 import onnxruntime
 import psutil
@@ -41,7 +40,7 @@
 def test_recognition_models(arch_name, input_shape, train_mode, mock_vocab):
     batch_size = 4
     reco_model = recognition.__dict__[arch_name](vocab=mock_vocab, pretrained=True, input_shape=input_shape)
-    assert isinstance(reco_model, keras.Model)
+    assert isinstance(reco_model, tf.keras.Model)
     input_tensor = tf.random.uniform(shape=[batch_size, *input_shape], minval=0, maxval=1)
     target = ["i", "am", "a", "jedi"]
 
@@ -195,7 +194,7 @@ def test_recognition_zoo_error():
 def test_models_onnx_export(arch_name, input_shape):
     # Model
     batch_size = 2
-    keras.backend.clear_session()
+    tf.keras.backend.clear_session()
     model = recognition.__dict__[arch_name](pretrained=True, exportable=True, input_shape=input_shape)
     # SAR, MASTER, ViTSTR export currently only available with constant batch size
     if arch_name in ["sar_resnet31", "master", "vitstr_small", "parseq"]:
diff --git a/tests/tensorflow/test_models_utils_tf.py b/tests/tensorflow/test_models_utils_tf.py
index b57b41b14b..4783a09b40 100644
--- a/tests/tensorflow/test_models_utils_tf.py
+++ b/tests/tensorflow/test_models_utils_tf.py
@@ -2,7 +2,7 @@
 
 import pytest
 import tensorflow as tf
-from keras.applications import ResNet50
+from tensorflow.keras.applications import ResNet50
 
 from doctr.models.classification import mobilenet_v3_small
 from doctr.models.utils import (