Update main.py

main.py

@@ -1,4 +1,5 @@
 from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import pickle
 import uvicorn