提交 3ed91642 authored 作者: Frederic's avatar Frederic

Fix crash with when split return only 1 output

上级 2ccf6add
......@@ -310,7 +310,8 @@ def local_gpu_split(node):
any([c != 'output' and isinstance(c.op, GpuFromHost) for c, idx
in outs_clients])):
new_op = GpuSplit(node.op.len_splits)
split_res = new_op(gpu_from_host(input), *node.inputs[1:])
split_res = new_op(gpu_from_host(input), *node.inputs[1:],
return_list=True)
return [host_from_gpu(o) for o in split_res]
return False
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论