Update keras import to fix build error.

`from keras import ops` or `import keras` are valid, but `import keras.ops` is no longer valid since ops is not a module in keras namespace, but points to the public API. PiperOrigin-RevId: 626326912
PAIR-code · Apr 19, 2024 · fbd8874 · fbd8874
1 parent e1f5756
commit fbd8874
Showing 1 changed file with 1 addition and 2 deletions.
diff --git 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