test quantized model

predict.py

@@ -8,6 +8,7 @@ import json
 import os
 import random
 from torchvision import transforms
+from torch.quantization import quantize_dynamic
 
 # Load labels
 with open("labels.json", "r") as f:
@@ -36,14 +37,25 @@ class SwinCustom(nn.Module):
         outputs = self.model(images)
         return outputs.logits
 
-model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="large_swin_best_model.pth")
+# Download quantized model weights
+model_path = hf_hub_download(repo_id="Noha90/AML_16", filename="swin_large_quantised.pth")
 print("Model path:", model_path)
+
+# Build the model 
 model = SwinCustom(model_name=MODEL_NAME, num_classes=40)
+
+# Quantize the model
+quantized_model = quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)
+
+# Load the quantized state dict
 state_dict = torch.load(model_path, map_location="cpu")
 if "model_state_dict" in state_dict:
     state_dict = state_dict["model_state_dict"]
-model.load_state_dict(state_dict, strict=False)
-model.eval()
+quantized_model.load_state_dict(state_dict, strict=False)
+quantized_model.eval()
+
+# Use quantized_model
+model = quantized_model  
 
 # Preprocessing
 transform = transforms.Compose([