p47068915
/
mindspore
forked from mindspore-Ecosystem/mindspore
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

!2387 dataset=False when run in CPU 

Merge pull request !2387 from wukesong/cpu-lenet-sink-true
This commit is contained in:
mindspore-ci-bot 2020-06-20 15:10:00 +08:00 committed by Gitee
parent aad494a6fc fafa0d73fc
commit 228a1b87f8
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
model_zoo/lenet/train.py
View File

@ -40,6 +40,9 @@ if __name__ == "__main__":
args = parser.parse_args() 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) context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target)
ds_train = create_dataset(os.path.join(args.data_path, "train"), ds_train = create_dataset(os.path.join(args.data_path, "train"),
cfg.batch_size, cfg.batch_size,