From 14c46d93f840ec2cb0a0ecbf566529a686eec651 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=AF=85?= Date: Sat, 29 Sep 2018 19:07:51 +0800 Subject: [PATCH] Use T to determine output type first --- mace/python/tools/memory_optimizer.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/mace/python/tools/memory_optimizer.py b/mace/python/tools/memory_optimizer.py index 10a0ab50..a4688aad 100644 --- a/mace/python/tools/memory_optimizer.py +++ b/mace/python/tools/memory_optimizer.py @@ -148,6 +148,9 @@ class MemoryOptimizer(object): mem_id = self.op_mem.get(op.input[0], -1) else: output_type = mace_pb2.DT_FLOAT + for arg in op.arg: + if arg.name == 'T': + output_type = arg.i if len(op.output_type) > i: output_type = op.output_type[i] op_mem_block = self.get_op_mem_block( -- GitLab