提交 e0bf5038 authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Avoid digging through nodes that have multiple clients.

上级 7e2b9fcd
...@@ -33,12 +33,11 @@ def grab_cpu_scalar(v, nd): ...@@ -33,12 +33,11 @@ def grab_cpu_scalar(v, nd):
return v.dimshuffle(()) return v.dimshuffle(())
def find_node(v, cls): def find_node(v, cls, ignore_clients=False):
# This digs through possibly redundant transfers to for the node # This digs through possibly redundant transfers to for the node
# that has the op class specified. # that has the op class specified.
if v.owner is not None: if v.owner is not None and (ignore_clients or len(v.clients) == 1):
if (isinstance(v.owner.op, cls) and if isinstance(v.owner.op, cls):
len(v.clients) == 1):
return v.owner return v.owner
elif (isinstance(v.owner.op, GpuFromHost) and elif (isinstance(v.owner.op, GpuFromHost) and
v.owner.inputs[0].owner is not None and v.owner.inputs[0].owner is not None and
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论