Implemented suggestions from Oliver

NVIDIA-Merlin · Mar 15, 2023 · 2c0e9a2 · 2c0e9a2
1 parent b503c4c
commit 2c0e9a2
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/merlin/models/tf/core/encoder.py b/merlin/models/tf/core/encoder.py
@@ -166,7 +166,7 @@ def batch_predict(
 
         return merlin.io.Dataset(predictions)
 
-    def call(self, inputs, targets=None, training=False, testing=False, **kwargs):
+    def call(self, inputs, *, targets=None, training=False, testing=False, **kwargs):
         inputs = self._prepare_features(inputs, targets=targets)
         if isinstance(inputs, tuple):
             inputs, targets = inputs

diff --git a/merlin/models/tf/models/base.py b/merlin/models/tf/models/base.py
@@ -146,7 +146,7 @@ def get_output_schema(export_path: str) -> Schema:
             col_schema = ColumnSchema(
                 output_name,
                 dtype=output_spec.dtype.as_numpy_dtype,
-                dims=((0, None), (shape[1], shape[1])),
+                dims=(None, shape[1]),
             )
 
         output_schema.column_schemas[output_name] = col_schema