From b8e02348e651a56faa6267060a5ab5a6ca9ee809 Mon Sep 17 00:00:00 2001
From: ziheng <ziheng@apache.org>
Date: Mon, 26 Jun 2017 18:46:00 -0700
Subject: [PATCH] Fix] Avoid Directly Pass Python Context Object (#201)
---
examples/graph_executor/python/tvm_graph/build.py | 2 +-
examples/graph_executor/src/graph_executor.cc | 4 +++-
2 files changed, 4 insertions(+), 2 deletions(-)
diff --git a/examples/graph_executor/python/tvm_graph/build.py b/examples/graph_executor/python/tvm_graph/build.py
index 6ba51853b..228713d86 100644
--- a/examples/graph_executor/python/tvm_graph/build.py
+++ b/examples/graph_executor/python/tvm_graph/build.py
@@ -29,7 +29,7 @@ def build(sym, target, shape, dtype="float32"):
def bind(g, ctx):
- m = _create_exec(g.handle, ctx)
+ m = _create_exec(g.handle, ctx.device_type, ctx.device_id)
return m
diff --git a/examples/graph_executor/src/graph_executor.cc b/examples/graph_executor/src/graph_executor.cc
index 868cfa42b..d649017cf 100644
--- a/examples/graph_executor/src/graph_executor.cc
+++ b/examples/graph_executor/src/graph_executor.cc
@@ -263,7 +263,9 @@ tvm::runtime::Module CreateExecutor(nnvm::Graph g, TVMContext ctx) {
TVM_REGISTER_GLOBAL("tvm_graph._create_executor")
.set_body([](TVMArgs args, TVMRetValue *rv) {
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];
*rv = CreateExecutor(g, ctx);
});
--
GitLab