提交 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):
node.inputs[0].owner and
isinstance(node.inputs[0].owner.op,
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()
......
......@@ -514,7 +514,8 @@ def local_gpua_softmaxwithbias(node):
@register_opt('fast_compile')
@op_lifter([theano.tensor.opt.Assert])
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')
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论