Upload 2 files

app.py

@@ -22,15 +22,17 @@ def load_quantized_model(checkpoint_path):
         tie_word_embeddings=True,
     )
     
-    # Set the quantization configuration for the embedding layer
-    model.embed_tokens.qconfig = quantization.float_qparams_weight_only_qconfig
+    # Set quantization config for ALL embedding layers
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Embedding):
+            module.qconfig = quantization.float_qparams_weight_only_qconfig
     
     # Apply static quantization to the rest of the model
     model.qconfig = quantization.default_qconfig
     model = quantization.prepare(model, inplace=False)
     model = quantization.convert(model, inplace=False)
     
-    # Load the quantized checkpoint
+    # Load the checkpoint
     checkpoint = torch.load(checkpoint_path, map_location="cpu")
     model.load_state_dict(checkpoint["model_state_dict"])