Update app.py

app.py

@@ -147,10 +147,10 @@ def run_model_pipeline(model, original_image, comparison_image, threshold=0.5):
         st.subheader("Feature Heatmaps")
         col3, col4 = st.columns(2)
         with col3:
-            fig1 = visualize_heatmap(model, img1_tensor)
+            fig1 = visualize_heatmap(model, img1_tensor.float().squeeze(0))
             st.pyplot(fig1)
         with col4:
-            fig2 = visualize_heatmap(model, img2_tensor)
+            fig2 = visualize_heatmap(model, img2_tensor.float().squeeze(0))
             st.pyplot(fig2)
 
     else: