Spaces:

geitta
/

flower_classification

Sleeping

geitta commited on Oct 20, 2025

Commit

0434edb

verified ·

1 Parent(s): 22bb2bd

bug fix

Files changed (1) hide show

model.py CHANGED Viewed

@@ -11,7 +11,7 @@ def create_vit_model(num_classes: int = 5):
   for param in model.parameters():
     param.requires_grad = False
   # change the classifier head to suit our problem
-  vit.heads = nn.Sequential(nn.Linear(in_features=768,
                                       out_features=5,
                                       bias=True))
   return model, transforms

   for param in model.parameters():
     param.requires_grad = False
   # change the classifier head to suit our problem
+  model.heads = nn.Sequential(nn.Linear(in_features=768,
                                       out_features=5,
                                       bias=True))
   return model, transforms