Sandbox run src/api.py

src/api.py

@@ -1,22 +1,23 @@
 import torch
-from cnn import CNN  # Importing CNN class from cnn.py
 from fastapi import FastAPI, File, UploadFile
 from PIL import Image
 from torchvision import transforms
 
+from cnn import CNN  # Importing CNN class from cnn.py
+
 # Load the model
 model = CNN()
 model.load_state_dict(torch.load("mnist_model.pth"))
 model.eval()
 
 # Transform used for preprocessing the image
-transform = transforms.Compose([
-    transforms.ToTensor(),
-    transforms.Normalize((0.5,), (0.5,))
-])
+transform = transforms.Compose(
+    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
+)
 
 app = FastAPI()
 
+
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     image = Image.open(file.file).convert("L")