From 038dfc839f1bd6127c971973fc87b686ee5bbd12 Mon Sep 17 00:00:00 2001 From: Jinyu Xie Date: Tue, 14 Jul 2020 02:51:54 -0400 Subject: [PATCH] Fix pytorch frontend prim::Constant issue (#6051) --- python/tvm/relay/frontend/pytorch.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/tvm/relay/frontend/pytorch.py b/python/tvm/relay/frontend/pytorch.py index 0e04ffdc3a23..57b64ace64ae 100644 --- a/python/tvm/relay/frontend/pytorch.py +++ b/python/tvm/relay/frontend/pytorch.py @@ -2172,6 +2172,8 @@ def _get_constant(node): return node.f(attr_name) elif ty in ["TensorType", "CompleteTensorType"]: tensor = node.t(attr_name) + if tensor.is_cuda: + tensor = tensor.cpu() if len(tensor.shape) == 0: # tensor(0.1) # TODO(t-vi): When is this needed? return tensor.item()