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

Handling tag 'cpu' in the middle of graph

上级 55ceca97
......@@ -319,7 +319,8 @@ class GraphToGPU(NavigatorOptimizer):
context_name = None
for i in [mapping[i] for i in node.inputs]:
if isinstance(i.type, GpuArrayType):
target = getattr(i.tag, 'target', None)
if (isinstance(i.type, GpuArrayType) and target != 'cpu'):
context_name = i.type.context_name
move_to_GPU = True
break
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论