Update script.py

script.py

@@ -28,24 +28,29 @@ class PytorchWorker:
                 "cuda:0" if torch.cuda.is_available() else "cpu")
             print(f"Using devide: {self.device}")
 
-            model = timm.create_model(model_name,
-                                      num_classes=number_of_categories,
-                                      pretrained=False)
+            from torchvision.models import vit_b_16, ViT_B_16_Weights
+
+            model = vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1)
+            
+            model.heads.head = torch.nn.Linear(model.heads.head.in_features,
+                                               number_of_categories)
+
+            model.load_state_dict(torch.load(model_path))
 
             # if not torch.cuda.is_available():
             #     model_ckpt = torch.load(model_path, map_location=torch.device("cpu"))
             # else:
             #     model_ckpt = torch.load(model_path)
 
-            model_ckpt = torch.load(model_path, map_location=self.device)
-            model.load_state_dict(model_ckpt)
+            # model_ckpt = torch.load(model_path, map_location=self.device)
+            # model.load_state_dict(model_ckpt)
 
             return model.to(self.device).eval()
 
         self.model = _load_model(model_name, model_path)
 
         self.transforms = T.Compose([
-            T.Resize((336, 336)),
+            T.Resize((384, 384)),
             T.ToTensor(),
             T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
         ])
@@ -79,7 +84,6 @@ def make_submission(test_metadata, model_path, model_name, output_csv_path="./su
 
         predictions.append(np.argmax(logits))
 
-
     test_metadata["class_id"] = predictions
 
     user_pred_df = test_metadata.drop_duplicates("observation_id", keep="first")
@@ -92,7 +96,7 @@ if __name__ == "__main__":
     with zipfile.ZipFile("/tmp/data/private_testset.zip", 'r') as zip_ref:
         zip_ref.extractall("/tmp/data")
 
-    MODEL_PATH = "pytorch_model.bin"
+    MODEL_PATH = "model.pth"
     MODEL_NAME = "hf-hub:timm/eva02_large_patch14_clip_336.merged2b_ft_inat21"
 
     metadata_file_path = "./SnakeCLEF2024_TestMetadata.csv"