add use_mkldnn flag to v2

8 years ago · f878c0c6ba
parent 92c293aca8
commit f878c0c6ba
1 changed files with 2 additions and 0 deletions
--- a/python/paddle/v2/init.py
+++ b/python/paddle/v2/init.py
@ -78,6 +78,8 @@ def init(**kwargs):

    if 'use_gpu' in kwargs:
        cp.g_command_config_args['use_gpu'] = kwargs['use_gpu']
+    if 'use_mkldnn' in kwargs:
+        cp.g_command_config_args['use_mkldnn'] = kwargs['use_mkldnn']
    assert 'parallel_nn' not in kwargs, ("currently 'parallel_nn' is not "
                                         "supported in v2 APIs.")