From 61cef300ced6b1a1c49e6e5b9df1b1a7d796ea32 Mon Sep 17 00:00:00 2001
From: Jiarui Li <jiarui_li@outlook.com>
Date: Fri, 10 Jul 2026 05:41:38 +0800
Subject: [PATCH] Harden exposure autoencoder resume loading

---
 train_exposure_autoencoder.py | 43 +++++++++++++++++++++++++++--------
 1 file changed, 33 insertions(+), 10 deletions(-)

diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py
index 2355b98..d68d887 100644
--- a/train_exposure_autoencoder.py
+++ b/train_exposure_autoencoder.py
@@ -492,6 +492,15 @@ def save_autoencoder_checkpoint(payload: dict, checkpoint_path: Path) -> None:
     tmp_path.replace(checkpoint_path)
 
 
+def torch_load_checkpoint(checkpoint_path: Path, map_location):
+    try:
+        return torch.load(
+            checkpoint_path, map_location=map_location, weights_only=False
+        )
+    except TypeError:
+        return torch.load(checkpoint_path, map_location=map_location)
+
+
 def load_resume_checkpoint(
     checkpoint_path: Path,
     model,
@@ -505,7 +514,7 @@ def load_resume_checkpoint(
     show_progress: bool,
     logger,
 ) -> tuple[int, float, int, list[dict]]:
-    checkpoint = torch.load(checkpoint_path, map_location=device)
+    checkpoint = torch_load_checkpoint(checkpoint_path, map_location=device)
     if "model_state_dict" not in checkpoint:
         raise KeyError(
             f"Checkpoint does not contain model_state_dict: {checkpoint_path}"
@@ -577,13 +586,13 @@ def load_resume_checkpoint(
     return start_epoch, best_loss, stale_epochs, history
 
 
-def main() -> None:
-    args = parse_args()
-    resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint)
-    if resume_checkpoint is not None:
-        apply_resume_config(args, resume_checkpoint)
-    validate_args(args)
-    device, rank, local_rank, world_size = init_distributed(args)
+def run_training(
+    args: argparse.Namespace,
+    device: torch.device,
+    rank: int,
+    local_rank: int,
+    world_size: int,
+) -> None:
     set_seed(args.seed + rank)
     configure_torch_for_training(device)
     run_dir, run_name = distributed_run_dir(args, rank, world_size)
@@ -763,8 +772,22 @@ def main() -> None:
             break
     logger.info(f"Best validation loss: {best_loss:.6f}")
     logger.info(f"Checkpoint: {run_dir / 'best.pt'}")
-    if dist.is_initialized():
-        dist.destroy_process_group()
+
+
+def main() -> None:
+    args = parse_args()
+    resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint)
+    if resume_checkpoint is not None:
+        apply_resume_config(args, resume_checkpoint)
+    validate_args(args)
+    init_done = False
+    try:
+        device, rank, local_rank, world_size = init_distributed(args)
+        init_done = dist.is_initialized()
+        run_training(args, device, rank, local_rank, world_size)
+    finally:
+        if init_done and dist.is_initialized():
+            dist.destroy_process_group()
 
 
 if __name__ == "__main__":