From 4c457f82d2b1079cfaad954207312af362ca2ff1 Mon Sep 17 00:00:00 2001 From: Rohit7824567 Date: Thu, 19 Mar 2026 21:10:39 +0530 Subject: [PATCH] fix: return processed_images with predictions to prevent mask misalignment During Alaska coastal satellite image batch inference, resuming from checkpoint caused predicted masks to be saved under wrong filenames because processed_images (subset) was zipped with valid_images (full list). Returning both ensures correct image-mask pairing across interrupted 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