diff --git a/mace/python/tools/memory_optimizer.py b/mace/python/tools/memory_optimizer.py
index 8841ba577c704687efb9761a5a1e65eab7f4cbda..8b6997a3611a2e315359742a951195a5fa5267f5 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: