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

make the image2neigs gpu optimization try to move only the version implemented on the gpu.

上级 c0dd1eb5
...@@ -654,7 +654,8 @@ def gpu_images2neibs(ten4, neib_shape, neib_step=None, mode='valid'): ...@@ -654,7 +654,8 @@ def gpu_images2neibs(ten4, neib_shape, neib_step=None, mode='valid'):
@local_optimizer() @local_optimizer()
def use_gpu_images2neibs(node): def use_gpu_images2neibs(node):
if type(node.op) is Images2Neibs: if (type(node.op) is Images2Neibs and
node.op.mode in ['valid', 'wrap_centered']):
return [host_from_gpu(gpu_images2neibs(gpu_from_host(node.inputs[0]), return [host_from_gpu(gpu_images2neibs(gpu_from_host(node.inputs[0]),
node.inputs[1], node.inputs[2], node.inputs[1], node.inputs[2],
mode=node.op.mode))] mode=node.op.mode))]
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论