From 9029d70d3314c6248b0497f2986ab64083fa69b0 Mon Sep 17 00:00:00 2001
From: "sweep-nightly[bot]"
 <131841235+sweep-nightly[bot]@users.noreply.github.com>
Date: Sat, 25 Nov 2023 07:17:38 +0000
Subject: [PATCH 1/2] feat: Updated src/main.py

---
 src/main.py | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/main.py b/src/main.py
index 243a31e..99aa9fc 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,12 +1,19 @@
-from PIL import Image
+"""
+This module contains the implementation of a neural network model using PyTorch. It includes the data loading and preprocessing step,
+definition of the neural network architecture, and the training process for the MNIST dataset. It also contains code to save the trained model.
+"""
+
+import numpy as np
 import torch
 import torch.nn as nn
 import torch.optim as optim
-from torchvision import datasets, transforms
+from PIL import Image
 from torch.utils.data import DataLoader
-import numpy as np
+from torchvision import datasets, transforms
 
 # Step 1: Load MNIST Data and Preprocess
+# The 'transform' variable defines a series of preprocessing steps to be applied to the MNIST dataset images.
+# It converts images to PyTorch tensors and normalizes them so that the pixel values are in the range [-1, 1].
 transform = transforms.Compose([
     transforms.ToTensor(),
     transforms.Normalize((0.5,), (0.5,))
@@ -17,6 +24,12 @@
 
 # Step 2: Define the PyTorch Model
 class Net(nn.Module):
+    """
+    The Net class defines a simple feed-forward neural network model with one hidden layer.
+    Methods:
+        __init__: Initializes the layers of the neural network.
+        forward: Computes the forward pass of the network given an input tensor.
+    """
     def __init__(self):
         super().__init__()
         self.fc1 = nn.Linear(28 * 28, 128)
@@ -31,8 +44,11 @@ def forward(self, x):
         return nn.functional.log_softmax(x, dim=1)
 
 # Step 3: Train the Model
+# 'model' is an instance of the Net class which represents the neural network model.
 model = Net()
+# 'optimizer' defines the optimization algorithm (Stochastic Gradient Descent) used to update model parameters.
 optimizer = optim.SGD(model.parameters(), lr=0.01)
+# 'criterion' is the loss function used to measure how well the model performs during training. Here, Negative Log Likelihood Loss is used.
 criterion = nn.NLLLoss()
 
 # Training loop

From 48a4c0051ca7557b597b86a5d336b14d6e2972bd Mon Sep 17 00:00:00 2001
From: "sweep-nightly[bot]"
 <131841235+sweep-nightly[bot]@users.noreply.github.com>
Date: Sat, 25 Nov 2023 07:20:18 +0000
Subject: [PATCH 2/2] feat: Updated src/api.py

---
 src/api.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/api.py b/src/api.py
index 36c257a..ef1715b 100644
--- a/src/api.py
+++ b/src/api.py
@@ -1,8 +1,8 @@
-from fastapi import FastAPI, UploadFile, File
-from PIL import Image
 import torch
-from torchvision import transforms
+from fastapi import FastAPI, File, UploadFile
 from main import Net  # Importing Net class from main.py
+from PIL import Image
+from torchvision import transforms
 
 # Load the model
 model = Net()