Loss is Nan

Thanks for your work. Well, it's occurs Loss is Nan when training stage that follows Usage step in ReadME，the error description:

`[2024-07-25 11:44:08 root](affinequant.py 159): INFO === Start quantize layer 3 ===
[2024-07-25 11:44:48 root](affinequant.py 282): INFO layer 3 iter 0 loss:0.012547658756375313 norm:0.005398863460868597 max memory_allocated 25510.10009765625
[2024-07-25 11:45:24 root](affinequant.py 282): INFO layer 3 iter 1 loss:0.010688988491892815 norm:0.002033534459769726 max memory_allocated 25510.10009765625
[2024-07-25 11:46:00 root](affinequant.py 282): INFO layer 3 iter 2 loss:0.010013626888394356 norm:0.0012934970436617732 max memory_allocated 25510.10009765625
[2024-07-25 11:46:35 root](affinequant.py 282): INFO layer 3 iter 3 loss:0.009677225723862648 norm:0.0008282792987301946 max memory_allocated 25510.10009765625
[2024-07-25 11:47:11 root](affinequant.py 282): INFO layer 3 iter 4 loss:0.009517742320895195 norm:0.00047518432256765664 max memory_allocated 25510.10009765625
[2024-07-25 11:47:47 root](affinequant.py 282): INFO layer 3 iter 5 loss:0.009451627731323242 norm:0.00032417610054835677 max memory_allocated 25510.10009765625
[2024-07-25 11:48:23 root](affinequant.py 282): INFO layer 3 iter 6 loss:0.009420475922524929 norm:0.0002503570867702365 max memory_allocated 25510.10009765625
[2024-07-25 11:48:58 root](affinequant.py 282): INFO layer 3 iter 7 loss:0.009402991272509098 norm:0.0002355567121412605 max memory_allocated 25510.10009765625
[2024-07-25 11:49:34 root](affinequant.py 282): INFO layer 3 iter 8 loss:0.009390492923557758 norm:0.00022407056530937552 max memory_allocated 25510.10009765625
[2024-07-25 11:50:10 root](affinequant.py 282): INFO layer 3 iter 9 loss:0.00937829539179802 norm:0.00022252841154113412 max memory_allocated 25510.10009765625
[2024-07-25 11:50:47 root](affinequant.py 282): INFO layer 3 iter 10 loss:0.00936979427933693 norm:0.00040410575456917286 max memory_allocated 25510.10009765625
[2024-07-25 11:51:23 root](affinequant.py 282): INFO layer 3 iter 11 loss:0.009373501874506474 norm:0.0007587362779304385 max memory_allocated 25510.10009765625
[2024-07-25 11:51:59 root](affinequant.py 282): INFO layer 3 iter 12 loss:0.009358054026961327 norm:0.0008263972704298794 max memory_allocated 25510.10009765625
[2024-07-25 11:52:35 root](affinequant.py 282): INFO layer 3 iter 13 loss:0.009706217795610428 norm:0.004431357607245445 max memory_allocated 25510.10009765625
[2024-07-25 11:53:12 root](affinequant.py 282): INFO layer 3 iter 14 loss:0.009572970680892467 norm:0.001220526173710823 max memory_allocated 25510.10009765625
[2024-07-25 11:53:48 root](affinequant.py 282): INFO layer 3 iter 15 loss:0.009499846026301384 norm:0.0015964835183694959 max memory_allocated 25510.10009765625
[2024-07-25 11:54:25 root](affinequant.py 282): INFO layer 3 iter 16 loss:0.009589510038495064 norm:0.0031670858152210712 max memory_allocated 25510.10009765625
[2024-07-25 11:55:02 root](affinequant.py 282): INFO layer 3 iter 17 loss:0.009425442665815353 norm:0.0016205854481086135 max memory_allocated 25510.10009765625
[2024-07-25 11:55:14 root](affinequant.py 272): INFO Loss is NAN, stopping training
> /data/users/wangh/code/AffineQuant/quantize/affinequant.py(275)affinequant()
-> loss_list.append(loss.data)
(Pdb) `

this is my bash file ,how to fix it?
`CUDA_VISIBLE_DEVICES=0 python main.py \
--model /data2/models/Llama-2/llama-2-7b-hf \
--epochs 20 --output_dir ./log/llama-7b-w4a4 \
--eval_ppl --wbits 4 --abits 4 --lwc --let --aug_loss --use_matrix --sf 0.1 \
--tasks hendrycksTest,piqa,arc_easy,arc_challenge,boolq,hellaswag,winogrande`

Start with the Usage steps in the README and reproduce the reported command from main.py. Then inspect quantize/affinequant.py around lines 159, 272, 275, and 282 to trace when the loss becomes NaN and how training stops. Done means the documented training run no longer reaches the reported NaN condition, with the relevant behavior verified by the available output or tests.

Written by the indexing model from the issue text.