From 84134639924185a1d43a5127c0e7f3c1570822fb Mon Sep 17 00:00:00 2001 From: Rohit7824567 Date: Thu, 19 Mar 2026 21:12:41 +0530 Subject: [PATCH] add weights_only=True to torch.load for secure model loading coastline U-Net model loading used torch.load without weights_only=True, which allows arbitrary code execution via pickle in PyTorch >= 2.0 and raises FutureWarning. Adding weights_only=True restricts deserialization to tensor data only, securing model checkpoint loading during GSoC inference runs. --- training_pipeline/predict.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/training_pipeline/predict.py b/training_pipeline/predict.py index 27886f7..37e34cd 100644 --- a/training_pipeline/predict.py +++ b/training_pipeline/predict.py @@ -168,7 +168,7 @@ def load_trained_model(model_path, device="cuda"): Loaded U-Net model """ model = UNet(n_channels=3, n_classes=1) - model.load_state_dict(torch.load(model_path, map_location=device)) + model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True)) model.to(device) return model