Fix num_labels=101 for Dinov2 checkpoint
Browse files- app/main.py +5 -1
app/main.py
CHANGED
|
@@ -25,7 +25,11 @@ print("🚀 Loading model and label mapping...")
|
|
| 25 |
MODEL_ID = "Arew99/dinov2-costum"
|
| 26 |
|
| 27 |
print("🚀 Loading model and label mapping...")
|
| 28 |
-
model = Dinov2ForImageClassification.from_pretrained(
|
|
|
|
|
|
|
|
|
|
|
|
|
| 29 |
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
|
| 30 |
|
| 31 |
model.eval()
|
|
|
|
| 25 |
MODEL_ID = "Arew99/dinov2-costum"
|
| 26 |
|
| 27 |
print("🚀 Loading model and label mapping...")
|
| 28 |
+
model = Dinov2ForImageClassification.from_pretrained(
|
| 29 |
+
MODEL_ID,
|
| 30 |
+
num_labels=101,
|
| 31 |
+
ignore_mismatched_sizes=True
|
| 32 |
+
)
|
| 33 |
processor = AutoImageProcessor.from_pretrained("facebook/dinov2-large")
|
| 34 |
|
| 35 |
model.eval()
|