Update app.py

app.py

@@ -10,6 +10,17 @@ class Prediction:
     prediction: TorchTensor
 
 app = fastapi.FastAPI(docs_url="/")
+from transformers import ViTForImageClassification
+
+# Define the number of classes in your custom dataset
+num_classes = 20
+
+# Initialize the ViTForImageClassification model
+model = ViTForImageClassification.from_pretrained(
+    'google/vit-base-patch16-224-in21k',
+    num_labels=num_classes  # Specify the number of classes
+)
+
 model = Model()
 model.load_state_dict(torch.load('model_state.pth', map_location='cpu'))
 # Define a function to preprocess the input image