diff --git a/src/ports/postgres/modules/deep_learning/input_data_preprocessor.py_in b/src/ports/postgres/modules/deep_learning/input_data_preprocessor.py_in
index 68d067b..21838a3 100644
--- a/src/ports/postgres/modules/deep_learning/input_data_preprocessor.py_in
+++ b/src/ports/postgres/modules/deep_learning/input_data_preprocessor.py_in
@@ -280,7 +280,7 @@
                 {self.schema_madlib}.agg_array_concat(ARRAY[{i}_norm::{float32}[]]) AS {i}
                 """.format(**locals()))
             shape_sql.append("""
-                ARRAY[count, {j}]::SMALLINT[] AS {i}_shape
+                ARRAY[count, {j}]::INTEGER[] AS {i}_shape
                 """.format(**locals()))
             bytea_sql.append("""
                 {self.schema_madlib}.array_to_bytea({i}) AS {i}
@@ -291,7 +291,7 @@
                 {self.schema_madlib}.agg_array_concat(ARRAY[{i}]) AS {i}
                 """.format(**locals()))
             shape_sql.append("""
-                ARRAY[count, {j}]::SMALLINT[] AS {i}_shape
+                ARRAY[count, {j}]::INTEGER[] AS {i}_shape
                 """.format(**locals()))
             bytea_sql.append("""
                 {self.schema_madlib}.array_to_bytea({i}) AS {i}
diff --git a/src/ports/postgres/modules/deep_learning/madlib_keras_helper.py_in b/src/ports/postgres/modules/deep_learning/madlib_keras_helper.py_in
index 243acd1..66641b1 100644
--- a/src/ports/postgres/modules/deep_learning/madlib_keras_helper.py_in
+++ b/src/ports/postgres/modules/deep_learning/madlib_keras_helper.py_in
@@ -132,7 +132,7 @@
 
     if is_platform_pg():
         res = plpy.execute(
-            """ SELECT {0}::SMALLINT[] AS shape
+            """ SELECT {0} AS shape
                 FROM {1}
             """.format(shape_col, table_name))
         images_per_seg = [sum(r['shape'][0] for r in res)]