def update_internal_dict()

in optimum/amd/brevitas/accelerate_utils.py [0:0]

• 8 lines of code
• 5 McCabe index (conditional complexity)

def update_internal_dict(module):
    prefix = module._hf_hook.weights_map.prefix
    for key in module.state_dict().keys():
        # It might happen that we call an quantization's inner modules, and this cause some parameters to be
        # already on meta device. This is not a problem for their value but we need to check here
        curr_device = (recurse_getattr(module, key + ".data")).device
        if str(curr_device) != "meta":
            module._hf_hook.weights_map.dataset.state_dict[prefix + key] = (
                recurse_getattr(module, key + ".data")
            ).cpu()