diff --git a/black_forest_streamlit.py b/black_forest_streamlit.py
index d91f503..89e9e16 100644
--- a/black_forest_streamlit.py
+++ b/black_forest_streamlit.py
@@ -381,7 +381,7 @@ else:
             model,
             gen_time,
         ) in st.session_state.generated_images:
-            st.image(filepath, caption=f"Prompt: {prompt_text}", use_column_width=True)
+            st.image(filepath, caption=f"Prompt: {prompt_text}", use_container_width=True)
             st.write(f"**Seed**: {seed}")
             st.write(f"**Model**: {model}")
             st.write(f"**Generated**: {gen_time.strftime('%Y-%m-%d %H:%M:%S')}")