1 files changed, 7 insertions, 7 deletions
diff --git a/src/mlia/nn/rewrite/library/fc_layer.py b/src/mlia/nn/rewrite/library/fc_layer.py
index 2480500..041ce85 100644
--- a/src/mlia/nn/rewrite/library/fc_layer.py
+++ b/src/mlia/nn/rewrite/library/fc_layer.py
@@ -1,18 +1,18 @@
-# SPDX-FileCopyrightText: Copyright 2023, Arm Limited and/or its affiliates.
+# SPDX-FileCopyrightText: Copyright 2023-2024, Arm Limited and/or its affiliates.
 # SPDX-License-Identifier: Apache-2.0
 """Example rewrite with one fully connected layer."""
 from typing import Any
 
-import tensorflow as tf
+from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
 
-def get_keras_model(input_shape: Any, output_shape: Any) -> tf.keras.Model:
+def get_keras_model(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for rewrite."""
-    model = tf.keras.Sequential(
+    model = keras.Sequential(
         (
-            tf.keras.layers.InputLayer(input_shape=input_shape),
-            tf.keras.layers.Reshape([-1]),
-            tf.keras.layers.Dense(output_shape),
+            keras.layers.InputLayer(input_shape=input_shape),
+            keras.layers.Reshape([-1]),
+            keras.layers.Dense(output_shape),
         )
     )
     return model