DumbledoreWiz commited on
Commit
85dcc91
1 Parent(s): cf14c94

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -2
app.py CHANGED
@@ -27,8 +27,10 @@ for path in [model_path, config_path, preprocessor_path]:
27
  logging.info(f"Found file: {path}")
28
 
29
  # Load the model and feature extractor using the local directory
30
- feature_extractor = ViTFeatureExtractor.from_pretrained(model_dir)
31
- model = ViTForImageClassification.from_pretrained(model_dir)
 
 
32
 
33
  # Ensure the model is in evaluation mode
34
  model.eval()
 
27
  logging.info(f"Found file: {path}")
28
 
29
  # Load the model and feature extractor using the local directory
30
+ model_id = "google/vit-base-patch16-224"
31
+
32
+ feature_extractor = ViTFeatureExtractor.from_pretrained(model_id)
33
+ model = ViTForImageClassification.from_pretrained(model_path)
34
 
35
  # Ensure the model is in evaluation mode
36
  model.eval()