From 9939bea3f818b09cb6da5751bd8367b0cdbb326a Mon Sep 17 00:00:00 2001 From: nickwalton Date: Tue, 22 Oct 2019 16:13:46 -0600 Subject: [PATCH] update --- generator/ctrl/training_utils/training.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/generator/ctrl/training_utils/training.py b/generator/ctrl/training_utils/training.py index da17aba..6df42f0 100644 --- a/generator/ctrl/training_utils/training.py +++ b/generator/ctrl/training_utils/training.py @@ -165,7 +165,8 @@ run_config = tf.contrib.tpu.RunConfig( input_partition_dims=[[1, 1], [1, 1]], per_host_input_for_training=3)) tf.logging.set_verbosity(tf.logging.INFO) -model.fit(input_fn(), epochs=1) +params = {"batch_size": 2} +model.fit(input_fn(params=params), epochs=1) # estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)