Remove post-training TensorBoard log charts

app.py

@@ -20,6 +20,8 @@ from tensorboard.backend.event_processing import event_accumulator
 from torch.utils.tensorboard import SummaryWriter
 import matplotlib.pyplot as plt
 from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
+import matplotlib
+matplotlib.use('Agg')
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -252,11 +254,14 @@ def main():
         st.success("Training complete!")
 
         # Show last confusion matrix in Streamlit
-        st.subheader("Confusion Matrix (Last Epoch)")
-        fig, ax = plt.subplots()
-        disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
-        disp.plot(ax=ax)
-        st.pyplot(fig)
+        if 'cm' in locals():
+            st.subheader("Confusion Matrix (Last Epoch)")
+            fig, ax = plt.subplots()
+            disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
+            disp.plot(ax=ax)
+            st.pyplot(fig)
+        else:
+            st.info("Confusion matrix not found.")
 
 if __name__ == "__main__":
     main()