spuuntries
commited on
Commit
·
4ba204c
1
Parent(s):
3197cb6
fix: fix model link
Browse files
app.py
CHANGED
|
@@ -29,7 +29,7 @@ config = ViTConfig.from_pretrained("google/vit-base-patch16-224")
|
|
| 29 |
config.num_labels = 2 # Binary classification
|
| 30 |
|
| 31 |
# Download the model file
|
| 32 |
-
model_url = "https://huggingface.co/spuun/yummy-paws/resolve/main/
|
| 33 |
model_path = "best_model.pth"
|
| 34 |
|
| 35 |
if not os.path.exists(model_path):
|
|
@@ -39,7 +39,9 @@ if not os.path.exists(model_path):
|
|
| 39 |
|
| 40 |
# Load the trained model
|
| 41 |
model = ViTForImageClassification.from_pretrained(
|
| 42 |
-
model_path,
|
|
|
|
|
|
|
| 43 |
)
|
| 44 |
model.classifier = nn.Linear(model.config.hidden_size, 2)
|
| 45 |
model.to(device)
|
|
|
|
| 29 |
config.num_labels = 2 # Binary classification
|
| 30 |
|
| 31 |
# Download the model file
|
| 32 |
+
model_url = "https://huggingface.co/spuun/yummy-paws/resolve/main/model.pth"
|
| 33 |
model_path = "best_model.pth"
|
| 34 |
|
| 35 |
if not os.path.exists(model_path):
|
|
|
|
| 39 |
|
| 40 |
# Load the trained model
|
| 41 |
model = ViTForImageClassification.from_pretrained(
|
| 42 |
+
model_path,
|
| 43 |
+
config=config,
|
| 44 |
+
ignore_mismatched_sizes=True, # weights_only=False
|
| 45 |
)
|
| 46 |
model.classifier = nn.Linear(model.config.hidden_size, 2)
|
| 47 |
model.to(device)
|