Update app.py

app.py

@@ -225,8 +225,28 @@ with gr.Blocks(
         outputs=[attention_img, attention_label, head_gallery]
     )
 
-    prev_btn.click(step_prev, inputs=step_slider, outputs=step_slider)
-    next_btn.click(step_next, inputs=[step_slider, max_step_state], outputs=step_slider)
+    # Previous button
+    prev_btn.click(
+        step_prev,
+        inputs=step_slider,
+        outputs=step_slider
+    ).then(
+        update_display,
+        inputs=[step_slider, mode, avg_state, heads_state, labels_state],
+        outputs=[attention_img, attention_label, head_gallery]
+    )
+    
+    # Next button
+    next_btn.click(
+        step_next,
+        inputs=[step_slider, max_step_state],
+        outputs=step_slider
+    ).then(
+        update_display,
+        inputs=[step_slider, mode, avg_state, heads_state, labels_state],
+        outputs=[attention_img, attention_label, head_gallery]
+    )
+    
 
 
     # Updates on Mode Change