Skip to content
Snippets Groups Projects
Commit 73775f45 authored by Yizhi Liu's avatar Yizhi Liu Committed by Tianqi Chen
Browse files

AlterOpLayout with tvm.target (#463)

* AlterOpLayout with tvm.target

* fix test
parent 343eb82c
No related branches found
No related tags found
No related merge requests found
......@@ -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)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment