diff --git a/main.py b/main.py
index 8677279..4d5f7bd 100644
--- a/main.py
+++ b/main.py
@@ -147,6 +147,8 @@ def get_step_size(self, algorithm):
             # Update the internal state with the new step size as the minimum of the current and previous step sizes
             self.step_size = min(step_size, self.step_size)
             
+            self.initial_step_size = self.step_size
+            
             if self.counter < self.steps:
                 self.counter += 1