1 files changed, 16 insertions, 3 deletions
diff --git a/src/mlia/nn/rewrite/library/fc_clustering_layer.py b/src/mlia/nn/rewrite/library/fc_clustering_layer.py
index 07c07ac..72931c0 100644
--- a/src/mlia/nn/rewrite/library/fc_clustering_layer.py
+++ b/src/mlia/nn/rewrite/library/fc_clustering_layer.py
@@ -3,11 +3,24 @@
 """Example rewrite with one fully connected clustered layer."""
 from typing import Any
 
+import tensorflow_model_optimization as tfmot
 from keras.api._v2 import keras  # Temporary workaround for now: MLIA-1107
 
-from mlia.nn.rewrite.library.fc_layer import get_keras_model
-
 
 def get_keras_model_clus(input_shape: Any, output_shape: Any) -> keras.Model:
     """Generate TensorFlow Lite model for clustering rewrite."""
-    return get_keras_model(input_shape, output_shape)
+    clustering_params = {
+        "number_of_clusters": 32,
+        "cluster_centroids_init": tfmot.clustering.keras.CentroidInitialization.LINEAR,
+    }
+    model = tfmot.clustering.keras.cluster_weights(
+        to_cluster=keras.Sequential(
+            [
+                keras.layers.InputLayer(input_shape=input_shape),
+                keras.layers.Flatten(),
+                keras.layers.Dense(units=output_shape),
+            ]
+        ),
+        **clustering_params
+    )
+    return model