Update app.py

app.py

@@ -35,12 +35,12 @@ class ModelPre(torch.nn.Module):
     def __init__(self):
         super().__init__()
         self.embedding = torch.nn.Sequential(
-            *list(models.convnext_small(weights=models.ConvNeXt_Small_Weights.IMAGENET1K_V1).children())[:-1],
+            *list(models.efficientnet_v2_m(weights=models.EfficientNet_V2_M_Weights.IMAGENET1K_V1).children())[:-1],
             torch.nn.Flatten(),
-            torch.nn.Linear(in_features=768, out_features=512),
-            torch.nn.ReLU(),
-            torch.nn.Linear(in_features=512, out_features=len(le.classes_) + 1),
+            torch.nn.Linear(in_features=1280,out_features=512),
+            torch.nn.Linear(in_features=512,out_features=len_classes),
         )
+        # Freeze all layers
 
     def forward(self, data):
         return self.embedding(data)