agentic-fix(try-2): test_accurac issue

Author: raselmandol

train_mnist.py

@@ -104,8 +104,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)
-        batch_correct = (outputs.argmax(1) == labels).sum().item()
-        correct += batch_correct
+            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):
@@ -126,8 +126,8 @@ def forward(self, x):
         inputs, labels = inputs.to(device), labels.to(device)
         outputs = model(inputs)
         test_loss += criterion(outputs, labels).item()
-    batch_correct = (outputs.argmax(1) == labels).sum().item()
-    correct += batch_correct
+        batch_correct = (outputs.argmax(1) == labels).sum().item()
+        correct += batch_correct
         batch_acc = batch_correct / len(inputs)
         test_accuracies.append(batch_acc)