From fafa0d73fc16d412e7e21b7e15cbe96e31ade979 Mon Sep 17 00:00:00 2001
From: wukesong <wukesong1@huawei.com>
Date: Sat, 20 Jun 2020 11:28:57 +0800
Subject: [PATCH] dataset sink is false when run in CPU

---
 model_zoo/lenet/train.py | 3 +++
 1 file changed, 3 insertions(+)

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