feat: Updated src/main.py

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.ERROR)
+
 # Step 1: Load MNIST Data and Preprocess
 transform = transforms.Compose([
  transforms.ToTensor(),
@@ -39,10 +42,13 @@ def forward(self, x):
 epochs = 3
 for epoch in range(epochs):
  for images, labels in trainloader:
- optimizer.zero_grad()
- output = model(images)
- loss = criterion(output, labels)
- loss.backward()
- optimizer.step()
+ try:
+ optimizer.zero_grad()
+ output = model(images)
+ loss = criterion(output, labels)
+ loss.backward()
+ optimizer.step()
+ except Exception as e:
+ logging.exception("Exception occurred during training: %s", e)
 
 torch.save(model.state_dict(), "mnist_model.pth")