Update app.py

app.py

@@ -6,10 +6,13 @@ import torch
 from transformers import BertForSequenceClassification, BertTokenizer
 import gradio as gr
 
-# Load the model architecture
+import torch
+from transformers import BertForSequenceClassification
+
+# Load the model architecture with the number of labels
 model = BertForSequenceClassification.from_pretrained('bert-base-uncased', num_labels=2)
 
-# Load the state dict, ensuring tensors are on the CPU
+# Load the state dict while mapping to CPU
 try:
     model.load_state_dict(torch.load('bert_model_complete.pth', map_location=torch.device('cpu')), strict=False)
 except Exception as e: