diff --git a/lit_nlp/examples/models/instrumented_keras_lms.py b/lit_nlp/examples/models/instrumented_keras_lms.py
index 44236311..f9b4cbcd 100644
--- a/lit_nlp/examples/models/instrumented_keras_lms.py
+++ b/lit_nlp/examples/models/instrumented_keras_lms.py
@@ -6,8 +6,7 @@
 from typing import Sequence
 
 from absl import logging
-import keras.backend
-import keras.ops
+import keras
 from lit_nlp.api import model as lit_model
 from lit_nlp.api import types as lit_types
 from lit_nlp.lib import utils as lit_utils