Update app.py

app.py

@@ -497,7 +497,7 @@ def main(input_text, max_new_tokens):
     "drop_rate": 0.3,      # Dropout rate
     "qkv_bias": False      # Query-key-value bias
     }).to(device)
-    model.load_state_dict(weights)
+    model.load_state_dict(weights['model_state_dict'])
     model.eval()
 
     context_size = model.pos_emb.weight.shape[0]