fix network wide and deep

5 years ago · f3ae371436
parent ca6da6751f
commit f3ae371436
1 changed files with 2 additions and 1 deletions
--- 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__":