diff --git a/benchmark/op_count.py b/benchmark/op_count.py
index 117a0b5..7a5ade7 100644
--- a/benchmark/op_count.py
+++ b/benchmark/op_count.py
@@ -8,8 +8,8 @@ def main():
 
     compiled_reset = sim._reset.lower(sim.data, sim.default_data, None).compile()
     compiled_step = sim._step.lower(sim.data, 1).compile()
-    op_count_reset = compiled_reset.cost_analysis()[0]["flops"]
-    op_count_step = compiled_step.cost_analysis()[0]["flops"]
+    op_count_reset = compiled_reset.cost_analysis()["flops"]
+    op_count_step = compiled_step.cost_analysis()["flops"]
     print(f"Op counts:\n Reset: {op_count_reset}\n Step: {op_count_step}")