Merge remote-tracking branch 'origin/master'
This commit is contained in:
commit
45cc0ce3c4
1 changed files with 5 additions and 1 deletions
|
@ -122,7 +122,11 @@ def load_model_weights(model, checkpoint_file, sd_model_hash):
|
|||
pl_sd = torch.load(checkpoint_file, map_location="cpu")
|
||||
if "global_step" in pl_sd:
|
||||
print(f"Global Step: {pl_sd['global_step']}")
|
||||
sd = pl_sd["state_dict"]
|
||||
|
||||
if "state_dict" in pl_sd:
|
||||
sd = pl_sd["state_dict"]
|
||||
else:
|
||||
sd = pl_sd
|
||||
|
||||
model.load_state_dict(sd, strict=False)
|
||||
|
||||
|
|
Loading…
Reference in a new issue