From 8ab4a0deca2d864dbeafac7abd7cbf2366d91a26 Mon Sep 17 00:00:00 2001 From: yiyanzhi_akane Date: Wed, 26 May 2021 10:03:18 +0800 Subject: [PATCH] enable graph_kernel while train inceptionv3 on gpu --- model_zoo/official/cv/inceptionv3/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/model_zoo/official/cv/inceptionv3/train.py b/model_zoo/official/cv/inceptionv3/train.py index 4fd48db1bfe..5bab2c63506 100644 --- a/model_zoo/official/cv/inceptionv3/train.py +++ b/model_zoo/official/cv/inceptionv3/train.py @@ -60,6 +60,9 @@ if __name__ == '__main__': cfg = CFG_DICT[args_opt.platform] create_dataset = DS_DICT[cfg.ds_type] + if args_opt.platform == "GPU": + context.set_context(enable_graph_kernel=True) + context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.platform, save_graphs=False) if os.getenv('DEVICE_ID', "not_set").isdigit(): context.set_context(device_id=int(os.getenv('DEVICE_ID')))