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

Fix test crash and better handling of node with multiple output

上级 1d266293
...@@ -113,6 +113,7 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re, ...@@ -113,6 +113,7 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re,
reuse_out = out reuse_out = out
pre_allocated.add(out) pre_allocated.add(out)
allocated.add(ins) allocated.add(ins)
break
elif ins in view_of: elif ins in view_of:
origin = view_of[ins] origin = view_of[ins]
if ins in viewed_by[origin]: if ins in viewed_by[origin]:
...@@ -122,7 +123,7 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re, ...@@ -122,7 +123,7 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re,
not isinstance(origin, theano.Constant)): not isinstance(origin, theano.Constant)):
# where gc # where gc
for i in range(idx + 1, len(order)): for i in range(idx + 1, len(order)):
if reuse_out: if reuse_out is not None:
break break
for out in order[i].outputs: for out in order[i].outputs:
if (getattr(out, 'ndim', None) == 0 and if (getattr(out, 'ndim', None) == 0 and
...@@ -131,8 +132,8 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re, ...@@ -131,8 +132,8 @@ def calculate_reallocate_info(order, fgraph, storage_map, compute_map_re,
reuse_out = out reuse_out = out
pre_allocated.add(out) pre_allocated.add(out)
allocated.add(ins) allocated.add(ins)
break
if reuse_out: if reuse_out is not None:
reallocated_info[ins] = [ins, reuse_out] reallocated_info[ins] = [ins, reuse_out]
return reallocated_info return reallocated_info
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论