Add CNN class for training MNIST model and import it into main.py

[sweep-nightly]

Description

This PR adds a new CNN class in a new file named cnn.py to handle training a Convolutional Neural Network (CNN) on the MNIST dataset. The CNN class is implemented using PyTorch and includes methods for defining the architecture, training the model, and saving the trained model.

The main.py file has been modified to import the CNN class and use it to train the model on the MNIST dataset. After training, the model is saved as mnist_model.pth.

Summary of Changes

• Created a new file src/cnn.py to implement the CNN class.
• Imported necessary modules in cnn.py for PyTorch and torchvision.
• Defined a new class CNN in cnn.py that inherits from torch.nn.Module.
• Implemented the __init__ method in CNN class to define the architecture of the CNN.
• Implemented the forward method in CNN class for the forward pass of the CNN.
• Implemented the train method in CNN class to train the model on the MNIST dataset.
• Implemented the save_model method in CNN class to save the trained model.
• Updated src/main.py to import the CNN class from cnn.py.
• Created an instance of the CNN class in main.py and trained the model on the MNIST dataset.
• Saved the trained model as mnist_model.pth in main.py.

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

[sweep-nightly]

Rollback Files For Sweep

• Rollback changes to src/main.py
• Rollback changes to src/cnn.py

[sweep-nightly]

Apply Sweep Rules to your PR?

• Apply: All docstrings and comments should be up to date.
• Apply: Code should be properly formatted and indented.
• Apply: Variable and function names should be descriptive and follow a consistent naming convention.
• Apply: Imports should be organized and grouped together.
• Apply: There should be no unused imports or variables.
• Apply: Code should be properly commented and include docstrings for functions and classes.

[sweep-nightly]

Apply Sweep Rules to your PR?

• Apply: All docstrings and comments should be up to date.
• Apply: Code should be properly formatted and indented.
• Apply: Variable and function names should be descriptive and follow a consistent naming convention.
• Apply: Imports should be organized and grouped together.
• Apply: There should be no unused imports or variables.
• Apply: Code should be properly commented and include docstrings for functions and classes.