Update app.py

app.py

@@ -12,7 +12,7 @@ with open("class_names.txt", "r") as f:
 
 ### Model and transforms preparation ###
 # Create model and transforms
-model, _, _, transforms = create_vit(output_shape=101, classes=class_names)
+model, _, _, transforms = create_vit(output_shape=len(class_names), classes=class_names)
 
 model = torch.compile(model)