Skip to content

Commit

Permalink
feat: Updated src/main.py
Browse files Browse the repository at this point in the history
  • Loading branch information
sweep-nightly[bot] authored Nov 26, 2023
1 parent 129646e commit a411ef7
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions src/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np

# Step 1: Load MNIST Data and Preprocess
# The MNIST dataset is loaded and preprocessed by transforming the images to tensors and normalizing them.
Expand All @@ -16,6 +15,8 @@
trainset = datasets.MNIST('.', download=True, train=True, transform=transform)
# Step 2: Define the PyTorch Model
class Net(nn.Module):
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
# 'trainloader' variable added back into the code to load the training data for the model
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(28 * 28, 128)
Expand All @@ -36,7 +37,8 @@ def forward(self, x):

# Training loop
epochs = 3
for epoch in range(epochs):
for _ in range(epochs):
# 'epoch' variable removed as it is not used within the loop body
for images, labels in trainloader:
optimizer.zero_grad()
output = model(images)
Expand Down

0 comments on commit a411ef7

Please sign in to comment.