def set_initialized_submodules(model, state_dict_keys):
    """
    Sets the `_is_hf_initialized` flag in all submodules of a given model when all its weights are in the loaded state
    dict.
    """
    for module_name, module in model.named_modules():
        loaded_keys = [k.replace(f"{module_name}.", "") for k in state_dict_keys if k.startswith(f"{module_name}.")]
        if len(set(module.state_dict().keys()) - set(loaded_keys)) == 0:
            module._is_hf_initialized = True

    Sets the `_is_hf_initialized` flag in all submodules of a given model when all its weights are in the loaded state
    dict.
    
dpkd/transformers/src/transformers/modeling_utils.py