提交 b0ce28cd authored 作者: Frederic Bastien's avatar Frederic Bastien

[OPT FAILURE] Fix gh-4131

上级 0babc678
......@@ -1023,11 +1023,11 @@ def local_gpu_flatten(node):
return [gpu_flatten(host_input.owner.inputs[0], outdim)(
as_cuda_ndarray_variable(host_input.owner.inputs[0]))]
if isinstance(node.op, tensor.Flatten):
x, shp= node.inputs
outdim = node.op.outdim
x = node.inputs
if x.owner and isinstance(x.owner.op, HostFromGpu):
outdim = node.op.outdim
gpu_x, = x.owner.inputs
return [host_from_gpu(gpu_flatten(host_input.owner.inputs[0], outdim)(gpu_x))]
return [host_from_gpu(gpu_flatten(gpu_x, outdim))]
return False
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论