fix(api): unwrap state dict from VAE

ssube · Feb 17, 2023 · 37b173d · 37b173d
1 parent c74d22a
commit 37b173d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/api/onnx_web/convert/utils.py b/api/onnx_web/convert/utils.py
@@ -221,7 +221,9 @@ def load_tensor(name: str, map_location=None):
                     "failed to load with Torch JIT, falling back to PyTorch", e
                 )
                 checkpoint = torch.load(name, map_location=map_location)
-
+                checkpoint = (
+                    checkpoint["state_dict"] if "state_dict" in checkpoint else checkpoint
+                )
     else:
         logger.debug("loading ckpt")
         checkpoint = torch.load(name, map_location=map_location)