From 8e87d573233c4b575a9dc08a44a8eb44bde0f314 Mon Sep 17 00:00:00 2001
From: Jiarui Li <jiarui_li@outlook.com>
Date: Sun, 12 Jul 2026 09:57:36 +0800
Subject: [PATCH] Fix exposure encoding checkpoint loading

---
 encode_exposure_cache.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/encode_exposure_cache.py b/encode_exposure_cache.py
index 0c15e92..5f849fb 100644
--- a/encode_exposure_cache.py
+++ b/encode_exposure_cache.py
@@ -67,7 +67,12 @@ def main() -> None:
     if output_path.exists() and not args.overwrite:
         raise FileExistsError(f"{output_path} exists; pass --overwrite")
 
-    checkpoint_data = torch.load(args.checkpoint, map_location="cpu")
+    try:
+        checkpoint_data = torch.load(
+            args.checkpoint, map_location="cpu", weights_only=False
+        )
+    except TypeError:
+        checkpoint_data = torch.load(args.checkpoint, map_location="cpu")
     model_cfg = checkpoint_data["model_config"]
     normalization = checkpoint_data["normalization"]
     encoder = TimesNetExposureEncoder(