Skip to content

Commit ad9b053

Browse files
committed
fix for avoiding Buffer.vload(...) case
1 parent 54c6864 commit ad9b053

File tree

1 file changed

+6
-1
lines changed

1 file changed

+6
-1
lines changed

python/tvm/script/parser.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
use for error reporting.
2222
"""
2323
# pylint: disable=invalid-name, inconsistent-return-statements, no-else-return
24+
import types
2425
import json
2526
import operator
2627
import inspect
@@ -580,10 +581,14 @@ def transform_Assign(self, node):
580581
arg_list = self.parse_arg_list(func, node.rhs)
581582
func.handle(node, self.context, arg_list, node.rhs.func_name.span)
582583
return self.parse_body(node)
583-
elif callable(func):
584+
elif isinstance(func, types.FunctionType):
584585
# Pattern 5
585586
args = [self.transform(arg) for arg in node.rhs.params]
586587
out = func(*args)
588+
589+
if len(node.lhs) == 1 and not isinstance(out, list):
590+
out = [out]
591+
587592
assert len(out) == len(node.lhs)
588593

589594
for var, value in zip(node.lhs, out):

0 commit comments

Comments
 (0)