Fix chief-0 timeout in tf_keras premade models distributed test by specifying steps_per_epoch when training on Dataset.

PiperOrigin-RevId: 937158025

Author: tensorflower-gardener

tf_keras/distribute/keras_premade_models_test.py

@@ -119,7 +119,9 @@ def test_linear_model(self, distribution, use_dataset_creator, data_fn):
                     inputs, output = get_numpy()
                     hist = model.fit(inputs, output, epochs=3)
                 else:
-                    hist = model.fit(get_dataset(), epochs=3)
+                    hist = model.fit(
+                        get_dataset(), epochs=3, steps_per_epoch=INPUT_SIZE
+                    )
                 self.assertLess(hist.history["loss"][2], 0.2)
 
     @tf.__internal__.distribute.combinations.generate(
@@ -162,7 +164,9 @@ def test_wide_deep_model(self, distribution, use_dataset_creator, data_fn):
                     inputs, output = get_numpy()
                     hist = wide_deep_model.fit(inputs, output, epochs=3)
                 else:
-                    hist = wide_deep_model.fit(get_dataset(), epochs=3)
+                    hist = wide_deep_model.fit(
+                        get_dataset(), epochs=3, steps_per_epoch=INPUT_SIZE
+                    )
             self.assertLess(hist.history["loss"][2], 0.2)