Fix: App issue

app.py

@@ -115,7 +115,15 @@ class GPT(nn.Module):
 def load_model():
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     model = GPT(GPTConfig())
-    model.load_state_dict(torch.load('nano_gpt_model.pt', map_location=device))
+    
+    # Load the state dict with weights_only=True and remove the _orig_mod prefix
+    state_dict = torch.load('nano_gpt_model.pt', map_location=device, weights_only=True)
+    new_state_dict = {}
+    for key in state_dict.keys():
+        new_key = key.replace('_orig_mod.', '')
+        new_state_dict[new_key] = state_dict[key]
+    
+    model.load_state_dict(new_state_dict)
     model.to(device)
     model.eval()
     return model, device