Update script.py

script.py

@@ -28,9 +28,9 @@ class PytorchWorker:
                 "cuda:0" if torch.cuda.is_available() else "cpu")
             print(f"Using devide: {self.device}")
 
-            from torchvision.models import vit_b_16, ViT_B_16_Weights
+            from torchvision.models import vit_b_16
 
-            model = vit_b_16(weights=ViT_B_16_Weights.IMAGENET1K_SWAG_E2E_V1)
+            model = vit_b_16(weights=None)
             
             model.heads.head = torch.nn.Linear(model.heads.head.in_features,
                                                number_of_categories)
@@ -52,7 +52,7 @@ class PytorchWorker:
         self.model = _load_model(model_name, model_path)
 
         self.transforms = T.Compose([
-            T.Resize((384, 384)),
+            T.Resize((224, 224)),
             T.ToTensor(),
             T.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
         ])