提交 a0350219 authored 作者: sentient07's avatar sentient07

Added condition on input nodes

上级 38e2e460
...@@ -881,6 +881,11 @@ def local_gpua_softmaxwithbias(node, context_name): ...@@ -881,6 +881,11 @@ def local_gpua_softmaxwithbias(node, context_name):
@register_opt('fast_compile') @register_opt('fast_compile')
@op_lifter([theano.tensor.opt.Assert]) @op_lifter([theano.tensor.opt.Assert])
def local_assert(node, context_name): def local_assert(node, context_name):
# Check if input nodes are already on the GPU
for i in node.inputs:
if isinstance(i.type, GpuArrayType):
return
return [host_from_gpu(node.op(as_gpuarray_variable(node.inputs[0], return [host_from_gpu(node.op(as_gpuarray_variable(node.inputs[0],
context_name), context_name),
*node.inputs[1:]))] *node.inputs[1:]))]
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论