diff --git a/examples/keras/preprocess.py b/examples/keras/preprocess.py
index 8d69acf..efb1545 100644
--- a/examples/keras/preprocess.py
+++ b/examples/keras/preprocess.py
@@ -31,7 +31,7 @@ class Preprocess(object):
             image = Image.open(local_file)
     
         image = ImageOps.grayscale(image).resize((28, 28))
-        return np.array([np.array(image).flatten()])
+        return np.array([np.array(image)])
 
     def postprocess(self, data: Any, state: dict, collect_custom_statistics_fn=None) -> dict:
         # post process the data returned from the model inference engine
diff --git a/examples/pytorch/preprocess.py b/examples/pytorch/preprocess.py
index 8d69acf..efb1545 100644
--- a/examples/pytorch/preprocess.py
+++ b/examples/pytorch/preprocess.py
@@ -31,7 +31,7 @@ class Preprocess(object):
             image = Image.open(local_file)
     
         image = ImageOps.grayscale(image).resize((28, 28))
-        return np.array([np.array(image).flatten()])
+        return np.array([np.array(image)])
 
     def postprocess(self, data: Any, state: dict, collect_custom_statistics_fn=None) -> dict:
         # post process the data returned from the model inference engine