From 65f9fc3f489b28fd2035abe01c7e8581cc893e7d Mon Sep 17 00:00:00 2001 From: looop5 Date: Mon, 31 May 2021 15:29:09 +0800 Subject: [PATCH] enable graph kernel when training dqn on GPU back-end --- model_zoo/official/rl/dqn/train.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/model_zoo/official/rl/dqn/train.py b/model_zoo/official/rl/dqn/train.py index 42f56df737c..435c960a171 100644 --- a/model_zoo/official/rl/dqn/train.py +++ b/model_zoo/official/rl/dqn/train.py @@ -34,6 +34,9 @@ set_seed(1) if __name__ == "__main__": context.set_context(mode=context.GRAPH_MODE, device_target=args.device_target) + if args.device_target == 'GPU': + # Enable graph kernel + context.set_context(enable_graph_kernel=True, graph_kernel_flags="--enable_parallel_fusion") env = gym.make('CartPole-v1') cfg.state_space_dim = env.observation_space.shape[0] cfg.action_space_dim = env.action_space.n