From e5c26b1258ef61bd4f7e7e50feff75a42b4cb902 Mon Sep 17 00:00:00 2001
From: liuqi <liuqi10@xiaomi.com>
Date: Tue, 19 Dec 2017 11:17:39 +0800
Subject: [PATCH] Fix bug of memory optimization when no output shape.

---
 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 8841ba57..8b6997a3 100644
--- a/mace/python/tools/memory_optimizer.py
+++ b/mace/python/tools/memory_optimizer.py
@@ -47,6 +47,9 @@ class MemoryOptimizer(object):
         # reuse mem
         mem_id = self.idle_mem.pop()
 
+      if not op.output_shape:
+        print('There is no output shape information to do memory optimization.')
+        return
       op.mem_id = mem_id
       self.op_mem[self._op_to_tensor(op)] = mem_id
       if mem_id not in self.mem_block:
-- 
GitLab