Commit
·
c2a4d06
1
Parent(s):
9f3a8d5
bug fix
Browse files
model.py
CHANGED
|
@@ -70,6 +70,9 @@ def create_vitbase_model(
|
|
| 70 |
num_classes=num_classes
|
| 71 |
)
|
| 72 |
|
|
|
|
|
|
|
|
|
|
| 73 |
# Load the trained weights
|
| 74 |
vitbase16_model = load_model(
|
| 75 |
model=vitbase16_model,
|
|
|
|
| 70 |
num_classes=num_classes
|
| 71 |
)
|
| 72 |
|
| 73 |
+
# Compile the model
|
| 74 |
+
vitbase16_model = torch.compile(vitbase16_model, backend="aot_eager")
|
| 75 |
+
|
| 76 |
# Load the trained weights
|
| 77 |
vitbase16_model = load_model(
|
| 78 |
model=vitbase16_model,
|