def METHOD_NAME(state_dict: Dict[str, Any], keep_non_lora=False):
    '''
    if keep_non_lora, also return non_lora state_dict
    '''
    state_dict_lora = OrderedDict()
    state_dict_non_lora = OrderedDict()
    for k, v in state_dict.items():
        if 'lora_A' in k or 'lora_B' in k:
            state_dict_lora[k] = v
        elif keep_non_lora:
            state_dict_non_lora[k] = v
    if keep_non_lora:
        return state_dict_lora, state_dict_non_lora
    else:
        return state_dict_lora, None