From 918c0b5ae1917eabb497eae953ef480c3739a07c Mon Sep 17 00:00:00 2001
From: "sweep-nightly[bot]"
 <131841235+sweep-nightly[bot]@users.noreply.github.com>
Date: Tue, 24 Oct 2023 02:45:35 +0000
Subject: [PATCH] feat: Updated src/main.py

---
 src/main.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/src/main.py b/src/main.py
index 243a31e..5211dcd 100644
--- a/src/main.py
+++ b/src/main.py
@@ -5,6 +5,7 @@
 from torchvision import datasets, transforms
 from torch.utils.data import DataLoader
 import numpy as np
+import logging
 
 # Step 1: Load MNIST Data and Preprocess
 transform = transforms.Compose([
@@ -12,6 +13,8 @@
     transforms.Normalize((0.5,), (0.5,))
 ])
 
+logging.basicConfig(filename='training.log', level=logging.INFO, format='%(asctime)s %(message)s')
+
 trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
 trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
 
@@ -44,5 +47,6 @@ def forward(self, x):
         loss = criterion(output, labels)
         loss.backward()
         optimizer.step()
+        logging.info('Epoch: %s, Loss: %s', epoch, loss.item())
 
 torch.save(model.state_dict(), "mnist_model.pth")
\ No newline at end of file