Update app.py

app.py

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