feat: Updated src/main.py

sweepai · Oct 24, 2023 · 1ff1425 · 1ff1425
1 parent 2ab11cb
commit 1ff1425
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/main.py b/src/main.py
@@ -1,3 +1,4 @@
+import logging
 from PIL import Image
 import torch
 import torch.nn as nn
@@ -6,6 +7,8 @@
 from torch.utils.data import DataLoader
 import numpy as np
 
+logging.basicConfig(filename='training.log', level=logging.INFO, format='%(asctime)s %(message)s')
+
 # Step 1: Load MNIST Data and Preprocess
 transform = transforms.Compose([
     transforms.ToTensor(),
@@ -42,6 +45,7 @@ def forward(self, x):
         optimizer.zero_grad()
         output = model(images)
         loss = criterion(output, labels)
+        logging.info('Epoch: %s, Loss: %s', epoch, loss.item())
         loss.backward()
         optimizer.step()