diff --git a/tests/e2e/test_lightning_ref/expected_lightning_base_1step.png b/tests/e2e/test_lightning_ref/expected_lightning_base_1step.png
index 9034f08..8352efb 100644
Binary files a/tests/e2e/test_lightning_ref/expected_lightning_base_1step.png and b/tests/e2e/test_lightning_ref/expected_lightning_base_1step.png differ
diff --git a/tests/e2e/test_lightning_ref/expected_lightning_base_4step.png b/tests/e2e/test_lightning_ref/expected_lightning_base_4step.png
index 63df54c..d359ace 100644
Binary files a/tests/e2e/test_lightning_ref/expected_lightning_base_4step.png and b/tests/e2e/test_lightning_ref/expected_lightning_base_4step.png differ
diff --git a/tests/e2e/test_lightning_ref/expected_lightning_lora_4step.png b/tests/e2e/test_lightning_ref/expected_lightning_lora_4step.png
index dea0965..631adb8 100644
Binary files a/tests/e2e/test_lightning_ref/expected_lightning_lora_4step.png and b/tests/e2e/test_lightning_ref/expected_lightning_lora_4step.png differ