Upload model files

app.py

@@ -5,12 +5,14 @@ from fastapi import FastAPI, UploadFile, File
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
 
-app = FastAPI()
+# Allowlist the required global for PyTorch model loading
+torch.serialization.add_safe_globals(['cloudpickle.cloudpickle._make_skeleton_class'])
 
-# Load the pre-trained model
-model_uri = "model.pth"
-model = torch.load(model_uri)
+app = FastAPI()
 
+# Load the pre-trained model from within the container
+model_uri = "/app/model.pth"
+model = torch.load(model_uri, weights_only=True)
 
 # Define input schema for JSON requests
 class ImageInput(BaseModel):