agentic-fix(try-1): test_accurac issue

Author: raselmandol

train_mnist.py

@@ -104,6 +104,19 @@ def forward(self, x):
 save_markdown_log("test_output.md", test_log)
 # Mirror test log into docs/
 save_markdown_log(os.path.join("docs", "test_output.md"), test_log)
+        batch_correct = (outputs.argmax(1) == labels).sum().item()
+        correct += batch_correct
+test_log += predict_and_log_samples(model, test_loader, device)
+sample_imgs = save_test_samples_images(model, test_loader, device, out_dir="images", num_samples=5)
+for i, (idx, label, pred, img_path) in enumerate(sample_imgs, start=1):
+    img_name = os.path.basename(img_path)
+    test_log += f"![Sample {i:02d}](images/{img_name}) True: {label}, Pred: {pred}\n"
+
+# Add a test accuracy image that always overwrites
+save_test_accuracy_image(test_accuracies, out_dir="images")
+test_log += "\n## Test Accuracy Visual\n"
+test_log += "![Test Accuracy](images/test_accuracy.png)\n"
+save_markdown_log("test_output.md", test_log)
 model.eval()
 test_loss = 0
 correct = 0
@@ -141,6 +154,8 @@ def forward(self, x):
 save_markdown_log("test_output.md", test_log)
 # Mirror test log into docs/
 save_markdown_log(os.path.join("docs", "test_output.md"), test_log)
+# Mirror test log into docs/
+save_markdown_log(os.path.join("docs", "test_output.md"), test_log)
 
 # Ensure images are available under docs/images for GitHub Pages
 images_src = "images"