diff --git a/ethosu/vela/live_range.py b/ethosu/vela/live_range.py
index 7ff1b28..0b94b19 100644
--- a/ethosu/vela/live_range.py
+++ b/ethosu/vela/live_range.py
@@ -138,15 +138,11 @@
         return live_range
 
     def update_endtime(self):
-        self.end_time = 0
-        for rng in self.ranges.values():
-            self.end_time = max(self.end_time, rng.end_time)
+        self.end_time = self.current_time
         return self.end_time + 1
 
     def get_temporal_memory_usage(self, target_mem_area):
-        if not self.end_time:
-            self.update_endtime()
-        usage = np.zeros(self.end_time, dtype=np.int32)
+        usage = np.zeros(self.update_endtime(), dtype=np.int32)
         for rng in self.ranges.values():
             if rng.mem_area == target_mem_area:
                 # End time is inclusive
diff --git a/ethosu/vela/scheduler.py b/ethosu/vela/scheduler.py
index a375a93..de10bad 100644
--- a/ethosu/vela/scheduler.py
+++ b/ethosu/vela/scheduler.py
@@ -861,10 +861,12 @@
             return max_sched
 
         # Extract the cascades
-        cascades = [cascade for cascade in schedule.cascades.values()]
-        for cascade_info in cascades:
-            # Remove existing cascade from schedule
-            del schedule.cascades[cascade_info.end]
+        cascades = schedule.cascades
+        # Remove existing cascade from schedule
+        schedule.cascades = {}
+        for cost in schedule.cost_map.values():
+            cost.cascade = 0
+        for cascade_info in cascades.values():
             # Optimize the sub-schedule in this cascade
             opt_sub_schedule = self.optimize_sub_schedule(cascade_info, schedule, max_template, sram_limit)
             # Update the sub-schedule Op and cascade costs to the full schedule