Function convert_model_ckpt

diffpack/util.py:148–167

(state_dict)

def convert_model_ckpt(state_dict):
    new_state_dict = {}
    for k, v in state_dict.items():
        for i in range(4):
            if k.startswith(f"model_list.{i}.sigma_embed_layer"):
                new_k = k.replace(f"model_list.{i}.sigma_embed_layer", f"sigma_embedding_list.{i}")
                new_state_dict[new_k] = v
            elif k.startswith(f"model_list.{i}.gearnet"):
                new_k = k.replace(f"model_list.{i}.gearnet", f"model_list.{i}")
                new_state_dict[new_k] = v
            elif k.startswith(f"model_list.{i}.torsion_mlp"):
                new_k = k.replace(f"model_list.{i}.torsion_mlp", f"torsion_mlp_list.{i}")
                new_state_dict[new_k] = v
            elif k.startswith(f"confidence_model."):
                new_k = k
                new_state_dict[new_k] = v
            elif k.startswith(f"mlp."):
                new_k = k
                new_state_dict[new_k] = v
    return new_state_dict