提交 7ba334ce authored 作者: Roy Xue's avatar Roy Xue

fix crash, check first if the var is in variable_shape

上级 c0201d15
...@@ -642,12 +642,13 @@ class ProfileStats(object): ...@@ -642,12 +642,13 @@ class ProfileStats(object):
fct_shapes[node.fgraph].setdefault(node, []) fct_shapes[node.fgraph].setdefault(node, [])
sum_dense = 0 sum_dense = 0
for out in node.outputs: for out in node.outputs:
sh = self.variable_shape[out] if out in self.variable_shape.keys():
if hasattr(out.type, 'get_size'): sh = self.variable_shape[out]
v = out.type.get_size(sh) if hasattr(out.type, 'get_size'):
sum_dense += v v = out.type.get_size(sh)
else: sum_dense += v
v = "Unknown" else:
v = "Unknown"
var_mem[out] = v var_mem[out] = v
fct_memory[node.fgraph][node].append(v) fct_memory[node.fgraph][node].append(v)
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论