Harden exposure autoencoder resume loading

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:
+def run_training(
-    args = parse_args()
+    args: argparse.Namespace,
-    resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint)
+    device: torch.device,
-    if resume_checkpoint is not None:
+    rank: int,
-        apply_resume_config(args, resume_checkpoint)
+    local_rank: int,
-    validate_args(args)
+    world_size: int,
-    device, rank, local_rank, world_size = init_distributed(args)
+) -> None:
     set_seed(args.seed + rank)
     configure_torch_for_training(device)
     run_dir, run_name = distributed_run_dir(args, rank, world_size)
@@ -763,7 +772,21 @@ 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():
+
+
+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()