feat: Updated src/main.py

sweepai · Nov 26, 2023 · fccd851 · fccd851
1 parent a411ef7
commit fccd851
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/main.py b/src/main.py
@@ -14,9 +14,9 @@
 
 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
+class Net(nn.Module):
     def __init__(self):
         super().__init__()
         self.fc1 = nn.Linear(28 * 28, 128)