Upload app.py

app.py

@@ -1,7 +1,8 @@
 import torch
+import torch.nn as nn  # Added missing import
 import torch.ao.quantization as quantization
 from transformers import AutoTokenizer
-from model import TransformerModel  # Replace with your model class
+from model import TransformerModel
 import gradio as gr
 
 # Load the tokenizer
@@ -24,7 +25,7 @@ def load_quantized_model(checkpoint_path):
     
     # Set quantization config for ALL embedding layers
     for name, module in model.named_modules():
-        if isinstance(module, nn.Embedding):
+        if isinstance(module, nn.Embedding):  # Now works because `nn` is imported
             module.qconfig = quantization.float_qparams_weight_only_qconfig
     
     # Apply static quantization to the rest of the model
@@ -39,7 +40,6 @@ def load_quantized_model(checkpoint_path):
     model.eval()
     return model
 
-import gradio as gr
 
 # Load the quantized model
 model = load_quantized_model("checkpoint_quantized.pt")