AlterOpLayout with tvm.target (dmlc#463)

* AlterOpLayout with tvm.target * fix test
padmaramesh · May 3, 2018 · 63bd093 · 63bd093
1 parent 60c685d
commit 63bd093
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/nnvm/compiler/build_module.py b/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)