Relaxed test metrics

PiperOrigin-RevId: 599743562
tensorflow · Jan 19, 2024 · 207b647 · 207b647
1 parent 23fb1c8
commit 207b647
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/tensorflow_decision_forests/keras/BUILD b/tensorflow_decision_forests/keras/BUILD
@@ -119,7 +119,7 @@ py_test(
         "@ydf//yggdrasil_decision_forests/test_data",
     ],
     python_version = "PY3",
-    shard_count = 30,
+    shard_count = 50,
     deps = [
         ":core",
         ":keras",

diff --git a/tensorflow_decision_forests/keras/keras_test.py b/tensorflow_decision_forests/keras/keras_test.py
@@ -2878,7 +2878,8 @@ def test_node_format_blob_sequence(self):
     """Test that the node format is BLOB_SEQUENCE if required."""
     dataset = adult_dataset()
     model = keras.RandomForestModel(
-        advanced_arguments=keras.AdvancedArguments(node_format="BLOB_SEQUENCE")
+        num_trees=10,
+        advanced_arguments=keras.AdvancedArguments(node_format="BLOB_SEQUENCE"),
     )
     ds = keras.pd_dataframe_to_tf_dataset(dataset.train, dataset.label)
     model.fit(ds)