Add tests using mocker to main.py

[sweep-nightly]

Description

This PR adds tests to the main.py file using a mocker. The tests cover the data loading and preprocessing steps, as well as the model definition. The purpose of these tests is to ensure that the code in main.py is working as expected.

Summary of Changes

• Created a new file named test_main.py in the src directory.
• Imported the necessary modules for testing, including pytest and unittest.mock.
• Imported the functions and classes from main.py that were needed for testing.
• Added test functions for the data loading and preprocessing steps. Used unittest.mock to create mocks for the MNIST dataset and the DataLoader. Asserted that the correct transformations were applied to the dataset and that the DataLoader was created with the correct parameters.
• Added a test function for the model definition. Used unittest.mock to create a mock for the PyTorch model. Asserted that the model was defined correctly.
• Ensured that the code followed the user's preferences for code style and documentation. Added docstrings for each test function explaining what it tests and how. Used descriptive names for test functions and variables. Organized imports at the top of the file and removed any unused imports.

Please review and merge this PR. Thank you!

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
• 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.