Update app.py

app.py

@@ -1,6 +1,21 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+import torch
+import torch.nn as nn
 import torch
 from torchvision import transforms
 
+class Prediction(BaseModel):
+    prediction: torch.Tensor
+
+app = FastAPI()
+
+# Load the PyTorch model
+model = torch.load("best_model-epoch=01-val_loss=3.00.ckpt")
+
+# Define a function to preprocess the input
+
+
 def preprocess_input(input):
     """Preprocess the input image for the PyTorch image classification model.
 
@@ -49,3 +64,8 @@ async def predict_endpoint(input: fastapi.File):
 
     # Return the prediction.
     return Prediction(prediction=predicted_class)
+
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)