fix: test_accuracy

Author: raselmandol

train_mnist.py

@@ -88,15 +88,34 @@ def forward(self, x):
 )
 
 # Test Loop
+test_accuracy = correct / len(test_loader.dataset)
+test_loss /= len(test_loader)
+test_log += f"- **Test Loss**: {test_loss:.4f}\n"
+test_log += f"- **Test Accuracy**: {test_accuracy:.4f}\n\n"
+test_log += "## Sample Predictions\n\n"
+test_log += "| Image Index | True Label | Predicted Label |\n"
+test_log += "|-------------|------------|------------------|\n"
+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)
+test_log += "\n### Sample Images\n"
+for i, (idx, label, pred, img_path) in enumerate(sample_imgs, start=1):
+test_log += "![Test Accuracy](images/test_accuracy.png)\n"
+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)
 model.eval()
 test_loss = 0
 correct = 0
+test_accuracies = []
 with torch.no_grad():
     for inputs, labels in test_loader:
         inputs, labels = inputs.to(device), labels.to(device)
         outputs = model(inputs)
         test_loss += criterion(outputs, labels).item()
-        correct += (outputs.argmax(1) == labels).sum().item()
+        batch_correct = (outputs.argmax(1) == labels).sum().item()
+        correct += batch_correct
+        batch_acc = batch_correct / len(inputs)
+        test_accuracies.append(batch_acc)
 
 test_accuracy = correct / len(test_loader.dataset)
 test_loss /= len(test_loader)
@@ -115,7 +134,7 @@ def forward(self, x):
     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_accuracy, out_dir="images")
+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)