diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py
index 1109a8f..31eb566 100644
--- a/recml/core/utils/keras_utils.py
+++ b/recml/core/utils/keras_utils.py
@@ -21,6 +21,7 @@
 from typing import Any
 
 from absl import logging
+
 import jax
 import keras
 import orbax.checkpoint as ocp