Refactor training loop from script to class

[sweep-nightly]

Description

This PR refactors the training loop from a script format to a class format for better organization and reusability. The training loop code is currently in the file src/main.py and has been moved into a new class named MNISTTrainer. The class encapsulates the training loop and related functionality, allowing for easy instantiation and use in other parts of the codebase.

Summary of Changes

• In src/main.py:

  • Created a new class named MNISTTrainer to encapsulate the training loop and related functionality.
  • Moved the code for loading and preprocessing the MNIST data into a method named load_data.
  • Moved the code for defining the PyTorch model into a method named define_model.
  • Created a new method named train_model to implement the training loop.
  • Created a new method named save_model to save the trained model to a file.
  • Updated docstrings and comments to reflect the changes.

• In src/api.py:

  • Imported the MNISTTrainer class from src/main.py.
  • Replaced the code for loading the model with an instantiation of the MNISTTrainer class and a call to its load_model method.
  • Updated docstrings and comments to reflect the changes.

These changes improve the organization and reusability of the training loop code, making it easier to maintain and use in other parts of the codebase.

Fixes #6.

---

🎉 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/api.py
• Rollback changes to src/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.

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