fix multi-server execution bug by adding sink size parameter.

5 years ago · f228e647c4
parent bf1ba1bdd2
commit f228e647c4
1 changed files with 1 additions and 1 deletions
--- a/model_zoo/official/cv/resnet/train.py
+++ b/model_zoo/official/cv/resnet/train.py
@ -188,4 +188,4 @@ if __name__ == '__main__':
    if args_opt.net == "se-resnet50":
        config.epoch_size = config.train_epoch_size
    model.train(config.epoch_size - config.pretrain_epoch_size, dataset, callbacks=cb,
-                dataset_sink_mode=(not args_opt.parameter_server))
+                sink_size=dataset.get_dataset_size(), dataset_sink_mode=(not args_opt.parameter_server))