quatize lead to shape error on Llama 2

Pulled from the last version of today, now since I'm trying to finetune with any kind of quantize I got this following error:

python finetune/lora.py --checkpoint_dir  checkpoints/meta-llama/Llama-2-13b-chat-hf  --data_dir  data/csv --out_dir out/lora/csv_llama2_13b_ticker_tunned_quantize --precision bf16-true --quantize bnb.nf4
{'eval_interval': 100, 'save_interval': 100, 'eval_iters': 100, 'eval_max_new_tokens': 100, 'log_interval': 1, 'devices': 1, 'learning_rate': 0.0003, 'batch_size': 128, 'micro_batch_size': 1, 'gradient_accumulation_iters': 128, 'max_iters': 100000, 'weight_decay': 0.01, 'lora_r': 8, 'lora_alpha': 16, 'lora_dropout': 0.05, 'lora_query': True, 'lora_key': False, 'lora_value': True, 'lora_projection': False, 'lora_mlp': False, 'lora_head': False, 'warmup_steps': 100}
Seed set to 1337
Loading model 'checkpoints/meta-llama/Llama-2-13b-chat-hf/lit_model.pth' with {'name': 'Llama-2-13b-chat-hf', 'hf_config': {'org': 'meta-llama', 'name': 'Llama-2-13b-chat-hf'}, 'block_size': 4096, 'vocab_size': 32000, 'padding_multiple': 64, 'padded_vocab_size': 32000, 'n_layer': 40, 'n_head': 40, 'n_embd': 5120, 'rotary_percentage': 1.0, 'parallel_residual': False, 'bias': False, 'lm_head_bias': False, 'n_query_groups': 40, 'shared_attention_norm': False, '_norm_class': 'RMSNorm', 'norm_eps': 1e-05, '_mlp_class': 'LLaMAMLP', 'gelu_approximate': 'none', 'intermediate_size': 13824, 'rope_condense_ratio': 1, 'rope_base': 10000, 'r': 8, 'alpha': 16, 'dropout': 0.05, 'to_query': True, 'to_key': False, 'to_value': True, 'to_projection': False, 'to_mlp': False, 'to_head': False, 'head_size': 128, 'rope_n_elem': 128}
Number of trainable parameters: 6,553,600
Traceback (most recent call last):
  File "/mnt/d/deep/lit-gpt/finetune/lora.py", line 387, in <module>
    CLI(setup)
  File "/mnt/d/deep/lit-gpt/.uvenv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 96, in CLI
    return _run_component(components, cfg_init)
  File "/mnt/d/deep/lit-gpt/.uvenv/lib/python3.10/site-packages/jsonargparse/_cli.py", line 181, in _run_component
    return component(**cfg)
  File "/mnt/d/deep/lit-gpt/finetune/lora.py", line 121, in setup
    fabric.launch(main, data_dir, checkpoint_dir, out_dir)
  File "/mnt/d/deep/lit-gpt/.uvenv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 834, in launch
    return self._wrap_and_launch(function, self, *args, **kwargs)
  File "/mnt/d/deep/lit-gpt/.uvenv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 920, in _wrap_and_launch
    return to_run(*args, **kwargs)
  File "/mnt/d/deep/lit-gpt/.uvenv/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 925, in _wrap_with_setup
    return to_run(*args, **kwargs)
  File "/mnt/d/deep/lit-gpt/finetune/lora.py", line 161, in main
    f"Number of non trainable parameters: {num_parameters(model, requires_grad=False):,}"
  File "/mnt/d/deep/lit-gpt/lit_gpt/utils.py", line 36, in num_parameters
    total += math.prod(p.quant_state.shape)
AttributeError: 'list' object has no attribute 'shape'

lightning @ git+https://github.com/Lightning-AI/lightning@532c723c8584903dc719458d0ad52861d51bc395
lightning-utilities==0.9.0
bitsandbytes==0.41.1  
numpy==1.26.1
xformers==0.0.22.post4
torch==2.1.0
torchmetrics==1.2.0
torchvision==0.16.0
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.18.1
nvidia-nvjitlink-cu12==12.3.52
nvidia-nvtx-cu12==12.1.105

Reproduce the command from the issue with the listed environment, then inspect finetune/lora.py around the call to num_parameters and lit_gpt/utils.py at line 36. Confirm that quantized Llama 2 fine-tuning no longer fails on the quant_state shape error and add or run the relevant validation if the repository provides one.