diff --git a/ethosu/vela/debug_database.py b/ethosu/vela/debug_database.py
index 77e13eb..6964808 100644
--- a/ethosu/vela/debug_database.py
+++ b/ethosu/vela/debug_database.py
@@ -23,7 +23,6 @@
 
 from . import numeric_util
 from .operation import Operation
-from .shape4d import Shape4D
 
 UntypedDict = Dict[Any, Any]
 UntypedList = List[Any]
@@ -79,18 +78,9 @@
                 src_uid = cls._sourceUID[parent]
             uid = len(cls._optimisedUID)
             cls._optimisedUID[op] = (uid, src_uid)
-            ofm_shape = op.ofm_shapes[0] if op.ofm_shapes else Shape4D(op.outputs[0].shape)
+            ofm_shape = numeric_util.full_shape(3, op.outputs[0].shape, 1)
             cls._optimisedTable.append(
-                [
-                    uid,
-                    src_uid,
-                    op.type,
-                    op.kernel.width,
-                    op.kernel.height,
-                    ofm_shape.width,
-                    ofm_shape.height,
-                    ofm_shape.depth,
-                ]
+                [uid, src_uid, op.type, op.kernel.width, op.kernel.height, ofm_shape[-2], ofm_shape[-3], ofm_shape[-1]]
             )
 
     @classmethod
diff --git a/ethosu/vela/graph_optimiser.py b/ethosu/vela/graph_optimiser.py
index 1128a31..0754f7e 100644
--- a/ethosu/vela/graph_optimiser.py
+++ b/ethosu/vela/graph_optimiser.py
@@ -276,6 +276,7 @@
             scaled_op.attrs["rescale"] = 1 / 128
         elif "rescale" in scaled_op.attrs:
             del scaled_op.attrs["rescale"]
+        scaled_op.set_ifm_ofm_shapes()
 
     return op
 
diff --git a/ethosu/vela/pass_packing.py b/ethosu/vela/pass_packing.py
index 0820034..35f5330 100644
--- a/ethosu/vela/pass_packing.py
+++ b/ethosu/vela/pass_packing.py
@@ -414,7 +414,7 @@
         else:
             ps.ifm_tensor = ifm_tensor
             ps.ifm2_tensor = None
-            if ps.primary_op is not None:
+            if ps.primary_op is not None and ps.primary_op.run_on_npu:
                 ps.ifm_shapes.append(ps.primary_op.ifm_shapes[0])
 
         ps.ofm_tensor = ofm_tensor