!18142 Enable graph_kernel while train NASNet on GPU

Merge pull request !18142 from hujiahui8/graph_kernel
This commit is contained in:
i-robot 2021-06-11 09:01:49 +08:00 committed by Gitee
commit 9c07a9c116
1 changed files with 1 additions and 1 deletions

View File

@ -48,7 +48,7 @@ if __name__ == '__main__':
if args_opt.platform != "GPU":
raise ValueError("Only supported GPU training.")
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.platform, save_graphs=False)
context.set_context(mode=context.GRAPH_MODE, device_target=args_opt.platform, enable_graph_kernel=True)
if os.getenv('DEVICE_ID', "not_set").isdigit():
context.set_context(device_id=int(os.getenv('DEVICE_ID')))