From 73775f45328711a57b9e0a09de3308ebcc875b3e Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Thu, 3 May 2018 08:55:04 -0700 Subject: [PATCH] AlterOpLayout with tvm.target (#463) * AlterOpLayout with tvm.target * fix test --- nnvm/python/nnvm/compiler/build_module.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/nnvm/python/nnvm/compiler/build_module.py b/nnvm/python/nnvm/compiler/build_module.py index 436b06c5..3080baf1 100644 --- a/nnvm/python/nnvm/compiler/build_module.py +++ b/nnvm/python/nnvm/compiler/build_module.py @@ -264,7 +264,8 @@ def build(graph, target=None, shape=None, dtype="float32", if _all_var_init: init_var = initialize_variables(shape, dtype) # Apply optimization - graph = optimize(graph, shape, dtype, layout) + with target: + graph = optimize(graph, shape, dtype, layout) # Precompute prune if params and cfg.pass_enabled("PrecomputePrune"): graph, params = precompute_prune(graph, params)