!2387 dataset=False when run in CPU

Merge pull request !2387 from wukesong/cpu-lenet-sink-true
2020-06-20 15:10:00 +08:00 · 2020-06-20 15:10:00 +08:00 · 228a1b87f8
parent aad494a6fc fafa0d73fc
commit 228a1b87f8
1 changed files with 3 additions and 0 deletions
--- a/model_zoo/lenet/train.py
+++ b/model_zoo/lenet/train.py
@ -40,6 +40,9 @@ if __name__ == "__main__":

    args = parser.parse_args()

+    if args.device_target == "CPU":
+        args.dataset_sink_mode = False
+
    context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
    ds_train = create_dataset(os.path.join(args.data_path, "train"),
                              cfg.batch_size,