提交 7b663007 authored 作者: James Bergstra's avatar James Bergstra

tiny optimization in opt

上级 642ffda6
...@@ -31,9 +31,9 @@ def register_opt(*tags, **kwargs): ...@@ -31,9 +31,9 @@ def register_opt(*tags, **kwargs):
@local_optimizer([]) @local_optimizer([])
def local_cut_gpu_host_gpu(node): def local_cut_gpu_host_gpu(node):
if tensor.opt.opt.check_chain(node, GpuFromHost(), HostFromGpu()): if tensor.opt.opt.check_chain(node, gpu_from_host, host_from_gpu):
return [node.inputs[0].owner.inputs[0]] return [node.inputs[0].owner.inputs[0]]
if tensor.opt.opt.check_chain(node, HostFromGpu(), GpuFromHost()): if tensor.opt.opt.check_chain(node, host_from_gpu, gpu_from_host):
return [node.inputs[0].owner.inputs[0]] return [node.inputs[0].owner.inputs[0]]
return False return False
gpu_cut_copies.register('cut_gpu_host_transfers', local_cut_gpu_host_gpu, 'fast_run', 'inplace', 'gpu') gpu_cut_copies.register('cut_gpu_host_transfers', local_cut_gpu_host_gpu, 'fast_run', 'inplace', 'gpu')
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论