diff --git a/ethosu/vela/compiler_driver.py b/ethosu/vela/compiler_driver.py
index a9e3839..cb47539 100644
--- a/ethosu/vela/compiler_driver.py
+++ b/ethosu/vela/compiler_driver.py
@@ -146,14 +146,14 @@
             )
 
 
-def compiler_driver(nng, arch, options, scheduler_options):
+def compiler_driver(nng, arch, options, scheduler_options, network_type):
     assert verify_graph_health(nng)
 
     # Pre-optimisation operator tracking
     for sg in nng.subgraphs:
         visit_graph_post_order(sg.output_tensors, arch, [], [_record_operator])
 
-    nng = graph_optimiser.optimise_graph_a(nng, arch, options.verbose_graph)
+    nng = graph_optimiser.optimise_graph(nng, arch, network_type, options.verbose_graph)
     assert verify_graph_health(nng)
 
     if options.verbose_quantization: