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__":