feat: Updated src/main.py

sweepai · Nov 26, 2023 · 129646e · 129646e
1 parent daebccc
commit 129646e
Showing 1 changed file with 0 additions and 8 deletions.
diff --git a/src/main.py b/src/main.py
@@ -36,7 +36,6 @@ def forward(self, x):
 
 # Training loop
 epochs = 3
-epochs = 3
 for epoch in range(epochs):
     for images, labels in trainloader:
         optimizer.zero_grad()
@@ -45,11 +44,4 @@ def forward(self, x):
         loss.backward()
         optimizer.step()
 
-torch.save(model.state_dict(), "mnist_model.pth")
-
-torch.save(model.state_dict(), "mnist_model.pth")
-        loss = criterion(output, labels)
-        loss.backward()
-        optimizer.step()
-
 torch.save(model.state_dict(), "mnist_model.pth")