diff --git a/scripts/keras_MnWE.py b/scripts/keras_MnWE.py
index 0788e43..88b920a 100644
--- a/scripts/keras_MnWE.py
+++ b/scripts/keras_MnWE.py
@@ -26,7 +26,7 @@ def make_model():
     dense = layers.Dense(64, activation="relu")
     x = dense(inputs)
     outputs = layers.Dense(10)(x)
-    model = keras.ServiceEstimator(inputs=inputs, outputs=outputs, name="mnist_model")
+    model = keras.Model(inputs=inputs, outputs=outputs, name="mnist_model")
     model.compile(
         loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
         optimizer=keras.optimizers.RMSprop(),