Commit b8e02348 by ziheng Committed by Tianqi Chen

Fix] Avoid Directly Pass Python Context Object (#201)

parent 3cad2478
...@@ -29,7 +29,7 @@ def build(sym, target, shape, dtype="float32"): ...@@ -29,7 +29,7 @@ def build(sym, target, shape, dtype="float32"):
def bind(g, ctx): def bind(g, ctx):
m = _create_exec(g.handle, ctx) m = _create_exec(g.handle, ctx.device_type, ctx.device_id)
return m return m
......
...@@ -263,7 +263,9 @@ tvm::runtime::Module CreateExecutor(nnvm::Graph g, TVMContext ctx) { ...@@ -263,7 +263,9 @@ tvm::runtime::Module CreateExecutor(nnvm::Graph g, TVMContext ctx) {
TVM_REGISTER_GLOBAL("tvm_graph._create_executor") TVM_REGISTER_GLOBAL("tvm_graph._create_executor")
.set_body([](TVMArgs args, TVMRetValue *rv) { .set_body([](TVMArgs args, TVMRetValue *rv) {
void* graph_handle = args[0]; void* graph_handle = args[0];
TVMContext ctx = args[1]; int device_type = args[1];
int device_id = args[2];
TVMContext ctx{static_cast<DLDeviceType>(device_type), device_id};
nnvm::Graph g = static_cast<nnvm::Graph*>(graph_handle)[0]; nnvm::Graph g = static_cast<nnvm::Graph*>(graph_handle)[0];
*rv = CreateExecutor(g, ctx); *rv = CreateExecutor(g, ctx);
}); });
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment