Add comments and docstrings to main.py and api.py

[sweep-nightly]

Description

This PR adds comments and docstrings to the main.py and api.py files in order to improve code readability and maintainability. The comments and docstrings provide explanations of the code's purpose, functionality, and important steps.

Summary of Changes

main.py

• Added a module-level docstring that explains the purpose of the file.
• Added a docstring to the Net class that explains its purpose and functionality.
• Added comments to the init method of the Net class that explain each layer of the neural network.
• Added a comment to the forward method of the Net class that explains the data processing steps.
• Added comments to the code for loading and preprocessing the MNIST data that explain each step.
• Added comments to the code for training the model that explain each step.

api.py

• Added a module-level docstring that explains the purpose of the file.
• Added a docstring to the predict function that explains its purpose, parameters, and return value.
• Added comments to the code for loading the model and defining the transform that explain each step.
• Added comments to the code inside the predict function that explain the image processing and prediction steps.

These changes improve the code's documentation and make it easier for developers to understand and maintain the codebase.

Fixes #11.

---

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