Update app.py

app.py

@@ -100,7 +100,7 @@ def predict_and_display(image_path, model):
         _, predicted = torch.max(outputs, 1)
 
     # Display the image and prediction
-    st.image(image, caption='Uploaded Image', use_column_width=True)
+    st.image(image, caption='Uploaded Image', use_container_width=True)
     st.write(f"Predicted Digit: {predicted.item()}")
 
 # Streamlit UI