Skip to content

Commit 0376da1

Browse files
committed
make it possible to save nai model using safetensors
1 parent bb11bee commit 0376da1

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

modules/sd_models.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -144,8 +144,8 @@ def transform_checkpoint_dict_key(k):
144144

145145

146146
def get_state_dict_from_checkpoint(pl_sd):
147-
if "state_dict" in pl_sd:
148-
pl_sd = pl_sd["state_dict"]
147+
pl_sd = pl_sd.pop("state_dict", pl_sd)
148+
pl_sd.pop("state_dict", None)
149149

150150
sd = {}
151151
for k, v in pl_sd.items():

0 commit comments

Comments
 (0)