Browse Source

!2387 dataset=False when run in CPU

Merge pull request !2387 from wukesong/cpu-lenet-sink-true
tags/v0.5.0-beta
mindspore-ci-bot Gitee 6 years ago
parent
commit
228a1b87f8
1 changed files with 3 additions and 0 deletions
  1. +3
    -0
      model_zoo/lenet/train.py

+ 3
- 0
model_zoo/lenet/train.py View File

@@ -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,


Loading…
Cancel
Save