use proper argument for FixedShape (#169)

Co-authored-by: Karl Higley <kmhigley@gmail.com>
NVIDIA-Merlin · Nov 18, 2022 · b298635 · b298635
1 parent 4c4e0f1
commit b298635
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/merlin/schema/io/tensorflow_metadata.py b/merlin/schema/io/tensorflow_metadata.py
@@ -20,7 +20,14 @@
 import numpy
 
 from merlin.schema.io import proto_utils, schema_bp
-from merlin.schema.io.schema_bp import Feature, FeatureType, FixedShape, FloatDomain, IntDomain
+from merlin.schema.io.schema_bp import (
+    Feature,
+    FeatureType,
+    FixedShape,
+    FixedShapeDim,
+    FloatDomain,
+    IntDomain,
+)
 from merlin.schema.io.schema_bp import Schema as ProtoSchema
 from merlin.schema.io.schema_bp import ValueCount
 from merlin.schema.schema import ColumnSchema
@@ -281,7 +288,7 @@ def _pb_feature(column_schema):
         max_length = value_count.get("max")
 
         if min_length and max_length and min_length == max_length:
-            feature.shape = FixedShape(min_length)
+            feature.shape = FixedShape([FixedShapeDim(size=min_length)])
         elif min_length and max_length and min_length < max_length:
             feature.value_count = ValueCount(min=min_length, max=max_length)
         else: