diff --git a/ethosu/vela/npu_performance.py b/ethosu/vela/npu_performance.py
index e09fc9e..fc148f3 100644
--- a/ethosu/vela/npu_performance.py
+++ b/ethosu/vela/npu_performance.py
@@ -41,10 +41,7 @@
 
     if ps2.npu_block_type in set((NpuBlockType.ConvolutionMxN, NpuBlockType.VectorProduct)):
         op = ps2.primary_op
-        ifm_idx, _, _, _, _ = op.get_ifm_ifm2_weight_bias_ofm_indices()
-        ifm_block_depth = arch.calc_ifm_block_depth(
-            op.inputs[ifm_idx].shape[-1], op.inputs[ifm_idx].dtype.size_in_bits()
-        )
+        ifm_block_depth = arch.calc_ifm_block_depth(op.ifm.shape[-1], op.ifm.dtype.size_in_bits())
     else:
         ifm_block_depth = block_config_ps2[-1]
 
@@ -237,8 +234,8 @@
     elif primary_op:
         skirt = primary_op.attrs.get("skirt", skirt)
         explicit_padding = primary_op.attrs.get("explicit_padding", explicit_padding)
-        assert primary_op.attrs["npu_block_type"] == ps.npu_block_type
-        npu_block_type = primary_op.attrs["npu_block_type"]
+        assert primary_op.type.npu_block_type == ps.npu_block_type
+        npu_block_type = primary_op.type.npu_block_type
 
         ifm_tensor, _, weight_tensor, ofm_tensor = ps.get_primary_op_ifm_ifm2_weights_ofm()