diff --git a/ethosu/vela/live_range.py b/ethosu/vela/live_range.py
index dbc0ce4..14e83a3 100644
--- a/ethosu/vela/live_range.py
+++ b/ethosu/vela/live_range.py
@@ -194,10 +194,7 @@
 
 
 def extract_live_ranges_from_passes(
-    sg,
-    target_mem_area,
-    target_mem_type_set=set((MemType.Scratch, MemType.Scratch_fast)),
-    ignore_subgraph_input_output_tensors=False,
+    sg, target_mem_area, target_mem_type_set=None, ignore_subgraph_input_output_tensors=False,
 ):
     lr_graph = LiveRangeGraph()
 
@@ -205,6 +202,9 @@
         lr_graph.ignore_tensors.update(sg.input_tensors)
         lr_graph.ignore_tensors.update(sg.output_tensors)
 
+    if target_mem_type_set is None:
+        target_mem_type_set = set((MemType.Scratch, MemType.Scratch_fast))
+
     # Try to merge live ranges of operations in the NPU subgraphs
     if sg.placement == PassPlacement.Npu:
         merge_op_ranges(sg, lr_graph, target_mem_area, target_mem_type_set)