Update app.py

app.py

@@ -51,7 +51,7 @@ model = CustomDialoGPT(vocab_size, n_embd, n_head, n_layer)
 # Download and load model weights
 try:
     pth_filepath = hf_hub_download(repo_id=model_repo, filename=model_filename)
-    checkpoint = torch.load(pth_filepath, map_location=device)
+    checkpoint = torch.load(pth_filepath, map_location=device, weights_only=True)
 
     # Handle different checkpoint saving formats if needed.
     # If your checkpoint is just the state_dict, load it directly.