diff --git a/ethosu/vela/register_command_stream_generator.py b/ethosu/vela/register_command_stream_generator.py
index 2043127..b4a633e 100644
--- a/ethosu/vela/register_command_stream_generator.py
+++ b/ethosu/vela/register_command_stream_generator.py
@@ -583,6 +583,7 @@
         block_config,
         arch,
         block_type,
+        npu_op.ofm.shape,
         ifm_shape,
         ifm2_shape,
         uses_scalar,