feat: Updated src/api.py

src/api.py

@@ -2,12 +2,13 @@
 from PIL import Image
 import torch
 from torchvision import transforms
-from main import Net # Importing Net class from main.py
-
-# Load the model
-model = Net()
-model.load_state_dict(torch.load("mnist_model.pth"))
-model.eval()
+from main import MNISTTrainer # Importing MNISTTrainer class from main.py
+trainer = MNISTTrainer()
+trainloader = trainer.load_data()
+model = trainer.define_model()
+trainer.train_model(trainloader, model)
+trainer.save_model(model)
+model.eval() # Set the model to evaluation mode
 
 # Transform used for preprocessing the image
 transform = transforms.Compose([