Update utils.py

This commit is contained in:
Stardust·减
2023-08-21 19:21:04 +08:00
committed by GitHub
parent a9e5be78d9
commit 276fa3a3b8

View File

@@ -36,7 +36,7 @@ def load_checkpoint(checkpoint_path, model, optimizer=None, skip_optimizer=False
new_state_dict = {} new_state_dict = {}
for k, v in state_dict.items(): for k, v in state_dict.items():
try: try:
# assert "dec" in k or "disc" in k #assert "emb_g" not in k
# print("load", k) # print("load", k)
new_state_dict[k] = saved_state_dict[k] new_state_dict[k] = saved_state_dict[k]
assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape) assert saved_state_dict[k].shape == v.shape, (saved_state_dict[k].shape, v.shape)