Update app.py

app.py

@@ -1,5 +1,4 @@
 from fastapi import FastAPI
-from pydantic import BaseModel
 import torch
 import torch.nn as nn
 import torch
@@ -7,25 +6,8 @@ from torchvision import transforms
 
 from typing import Any, Type
 
-import pydantic
 import torch
 
-class TensorSchema(pydantic.BaseModel):
-  shape: list[int]
-  dtype: str
-  requires_grad: bool
-
-  def __get_pydantic_core_schema__(cls: Type[Any]) -> pydantic.schema.Schema:
-    return pydantic.schema.Schema(
-      type="object",
-      properties={
-        "shape": pydantic.schema.Schema(type="array", items=pydantic.schema.Schema(type="integer")),
-        "dtype": pydantic.schema.Schema(type="string"),
-        "requires_grad": pydantic.schema.Schema(type="boolean"),
-      },
-      required=["shape", "dtype", "requires_grad"],
-    )
-
 class TorchTensor(torch.Tensor):
   pass
 
@@ -87,7 +69,7 @@ async def predict_endpoint(input: fastapi.File):
   predicted_class = prediction.argmax(1)
 
   # Return the prediction.
-  return Prediction(prediction=predicted_class)
+  return {"prediction": predicted_class.item()}
 
 
 if __name__ == "__main__":