Add tests using mocker to main.py

[sweep-nightly]

Description

This PR adds tests using mocker to the main.py file in order to ensure the correctness of the MNIST digit classification pipeline. The tests cover data loading and preprocessing, model definition, and training. The tests are written in a separate test file, following best practices for test organization. The code is properly formatted, indented, and includes descriptive variable and function names, organized imports, and no unused imports or variables. The code is also properly commented and includes docstrings for functions and classes.

Summary of Changes

• Created a new test file src/test_main.py to contain the tests.
• Imported necessary modules such as torch, torchvision, pytest, and mocker.
• Imported the Net class and other necessary functions from main.py.
• Added a function to test data loading and preprocessing, using mocker to mock the data loading and transformation functions.
• Added a function to test the model definition, using mocker to mock the nn.Linear and nn.functional.relu functions.
• Added a function to test the training process, using mocker to mock the DataLoader, optimizer, and loss function.
• Added docstrings and comments to each function, explaining what the function does and what it tests.
• Formatted and indented the code properly, and ensured that variable and function names are descriptive and follow a consistent naming convention.
• Organized imports at the top of the file, and ensured there are no unused imports or variables.

Fixes #10.

---

🎉 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/test_main.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.