From bd9157f533e42ecd09d14ee2d1bda795e6a5b515 Mon Sep 17 00:00:00 2001 From: liuqi Date: Fri, 16 Mar 2018 14:45:33 +0800 Subject: [PATCH] Fix memory optimizer matrix type bug. --- mace/python/tools/memory_optimizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mace/python/tools/memory_optimizer.py b/mace/python/tools/memory_optimizer.py index 8753b270..e632a22a 100644 --- a/mace/python/tools/memory_optimizer.py +++ b/mace/python/tools/memory_optimizer.py @@ -35,7 +35,7 @@ class MemoryOptimizer(object): def get_mem_size(self, op_type, output_shape): mem_size = [0, 0] - if op_type == 'WinogradTransform' or op_type == 'GEMM': + if op_type == 'WinogradTransform' or op_type == 'MatMul': mem_size[0] = output_shape[2] * output_shape[3] mem_size[1] = output_shape[0] * int((output_shape[1]+3)/4) else: -- GitLab