From c3a6103af54a43a010f375a8a56eac6489d79fb8 Mon Sep 17 00:00:00 2001
From: City <125218114+city96@users.noreply.github.com>
Date: Fri, 7 Mar 2025 22:36:51 +0100
Subject: [PATCH] hyvid subkey load fix

---
 tools/convert.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/tools/convert.py b/tools/convert.py
index 001cd86..6a90960 100644
--- a/tools/convert.py
+++ b/tools/convert.py
@@ -144,7 +144,10 @@ def parse_args():
 def load_state_dict(path):
     if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]):
         state_dict = torch.load(path, map_location="cpu", weights_only=True)
-        state_dict = state_dict.get("model", state_dict)
+        for subkey in ["model", "module"]:
+            if subkey in state_dict:
+                state_dict = state_dict[subkey]
+                break
         if len(state_dict) < 20:
             raise RuntimeError(f"pt subkey load failed: {state_dict.keys()}")
     else: