Show only the last confusion matrix image after training

app.py

@@ -241,18 +241,18 @@ def main():
             y_pred[np.random.choice(100, 10, replace=False)] = np.random.randint(0, num_classes, 10)
             cm = confusion_matrix(y_true, y_pred, labels=range(num_classes))
 
-            # Plot and log confusion matrix as image
-            fig, ax = plt.subplots()
-            disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
-            disp.plot(ax=ax)
-            plt.close(fig)
-            writer.add_figure("confusion_matrix", fig, epoch)
+            # Only log confusion matrix in the last epoch
+            if epoch == epochs - 1:
+                fig, ax = plt.subplots()
+                disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=[f"Class {i}" for i in range(num_classes)])
+                disp.plot(ax=ax)
+                plt.close(fig)
+                writer.add_figure("confusion_matrix", fig, epoch)
 
             time.sleep(0.5)
 
         writer.close()
         st.success("Training complete!")
-        st.line_chart({"Loss": losses, "Val Loss": val_losses, "Accuracy": accuracies})
 
         # Wait a moment to ensure logs are written
         time.sleep(1)