This PR adds a new CNN class in a new file named cnn.py to handle the MNIST dataset. The CNN class is implemented using PyTorch and includes the model definition and training code. The trained model is then imported into main.py for further use.
Summary of Changes
Created a new file src/cnn.py to contain the CNN class.
Imported necessary PyTorch modules in cnn.py.
Defined a new class CNN in cnn.py that inherits from torch.nn.Module.
Implemented the __init__ method in CNN class to define the layers of the CNN.
Implemented the forward method in CNN class to perform the forward pass of the CNN.
Defined a function train_model in cnn.py to handle the training loop.
Modified src/main.py to import the CNN class and train_model function from cnn.py.
Instantiated the CNN class in main.py to create the model.
Defined a loss function and an optimizer in main.py.
Called the train_model function in main.py to train the model.
Please review and merge this PR to incorporate the changes.
Fixes #9.
🎉 Latest improvements to Sweep:
Sweep can now passively improve your repository! Check out Rules to learn more.
💡 To get Sweep to edit this pull request, you can:
Comment below, and Sweep can edit the entire PR
Comment on a file, Sweep will only modify the commented file
Edit the original issue to get Sweep to recreate the PR from scratch
Description
This PR adds a new CNN class in a new file named
cnn.py
to handle the MNIST dataset. The CNN class is implemented using PyTorch and includes the model definition and training code. The trained model is then imported intomain.py
for further use.Summary of Changes
src/cnn.py
to contain the CNN class.cnn.py
.CNN
incnn.py
that inherits fromtorch.nn.Module
.__init__
method inCNN
class to define the layers of the CNN.forward
method inCNN
class to perform the forward pass of the CNN.train_model
incnn.py
to handle the training loop.src/main.py
to import theCNN
class andtrain_model
function fromcnn.py
.CNN
class inmain.py
to create the model.main.py
.train_model
function inmain.py
to train the model.Please review and merge this PR to incorporate the changes.
Fixes #9.
🎉 Latest improvements to Sweep:
💡 To get Sweep to edit this pull request, you can: