diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 50368b8..f5006c6 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -109,6 +109,7 @@
         DebugDatabase.add_optimised(op, avgpool_op)
         avgpool_op.ifm_shapes.append(op.ifm_shapes[idx])
         avgpool_op.ofm_shapes.append(op.ofm_shapes[0])
+        avgpool_op.memory_function = Op.ConcatSliceWrite
     assert ofm.shape[axis] == offset
 
     # If axis corresponds to C-dimension, NHCWB16 can only be used in the output if all the concat_start's are a