From f3ae371436aadafb84f3814cfc7441a0b7f437b9 Mon Sep 17 00:00:00 2001 From: anzhengqi Date: Wed, 22 Jul 2020 15:30:44 +0800 Subject: [PATCH] fix network wide and deep --- .../recommend/wide_and_deep/train_and_eval_distribute.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/model_zoo/official/recommend/wide_and_deep/train_and_eval_distribute.py b/model_zoo/official/recommend/wide_and_deep/train_and_eval_distribute.py index 86317b66d5..5a7cf8c718 100644 --- a/model_zoo/official/recommend/wide_and_deep/train_and_eval_distribute.py +++ b/model_zoo/official/recommend/wide_and_deep/train_and_eval_distribute.py @@ -110,7 +110,8 @@ def train_and_eval(config): out = model.eval(ds_eval) print("=====" * 5 + "model.eval() initialized: {}".format(out)) model.train(epochs, ds_train, - callbacks=[TimeMonitor(ds_train.get_dataset_size()), eval_callback, callback, ckpoint_cb]) + callbacks=[TimeMonitor(ds_train.get_dataset_size()), eval_callback, callback, ckpoint_cb], + sink_size=ds_train.get_dataset_size()) if __name__ == "__main__":