diff --git a/model_zoo/lenet/train.py b/model_zoo/lenet/train.py index 80f30310ab1c6ffd8e3a5c652a39546452f3443c..740b6e8ca3597fe5404d171ff569ca7233e81897 100644 --- 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,