diff --git a/deepof/train_utils.py b/deepof/train_utils.py
index 6f227b1d14b2a645c20d18183527cf8c5252f5f0..e12a9e5c91a3a953bdfab6c08d3572a1b2bf477c 100644
--- a/deepof/train_utils.py
+++ b/deepof/train_utils.py
@@ -143,7 +143,7 @@ def tensorboard_metric_logging(
     hpms: Any,
     X: np.ndarray,
     y: np.ndarray,
-    model: tensorflow.python.keras.engine.functional.Functional,
+    model: tf.python.keras.engine.functional.Functional,
     predictor: float,
     pheno_class: float,
     rec: str,