提交 db4f55c5 authored 作者: Frederic's avatar Frederic

[small BUG] We where loosing the condition when moving Assert op to the GPU

上级 1edc80df
...@@ -1846,7 +1846,8 @@ def local_assert(node): ...@@ -1846,7 +1846,8 @@ def local_assert(node):
node.inputs[0].owner and node.inputs[0].owner and
isinstance(node.inputs[0].owner.op, isinstance(node.inputs[0].owner.op,
HostFromGpu)): HostFromGpu)):
return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0]))] return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0],
*node.inputs[1:]))]
@register_opt() @register_opt()
......
...@@ -514,7 +514,8 @@ def local_gpua_softmaxwithbias(node): ...@@ -514,7 +514,8 @@ def local_gpua_softmaxwithbias(node):
@register_opt('fast_compile') @register_opt('fast_compile')
@op_lifter([theano.tensor.opt.Assert]) @op_lifter([theano.tensor.opt.Assert])
def local_assert(node): def local_assert(node):
return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0]))] return [host_from_gpu(node.op(node.inputs[0].owner.inputs[0],
*node.inputs[1:]))]
@register_opt('fast_compile') @register_opt('fast_compile')
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论