Wrap SequencePredictLast with tf.function (#1001)

NVIDIA-Merlin · Feb 27, 2023 · 5d0e090 · 5d0e090
1 parent 63f5ee1
commit 5d0e090
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/merlin/models/tf/transforms/sequence.py b/merlin/models/tf/transforms/sequence.py
@@ -269,6 +269,7 @@ class SequencePredictLast(SequenceTransform):
         so that the tensors sequences can be processed
     """
 
+    @tf.function
     def call(
         self, inputs: TabularData, targets=None, training=False, testing=False, **kwargs
     ) -> Tuple:

diff --git a/tests/unit/tf/blocks/test_mlp.py b/tests/unit/tf/blocks/test_mlp.py
@@ -103,10 +103,10 @@ def test_mlp_model_with_sequential_features_and_combiner(
         model, loader, run_eagerly=run_eagerly, reload_model=True, fit_kwargs={"pre": predict_last}
     )
 
-    metrics = model.evaluate(loader, batch_size=8, steps=1, return_dict=True, pre=predict_last)
+    metrics = model.evaluate(loader, steps=1, return_dict=True, pre=predict_last)
     assert len(metrics) > 0
 
-    predictions = model.predict(loader, batch_size=8, steps=1)
+    predictions = model.predict(loader, steps=1)
     assert predictions.shape == (8, 51997)