CUDA driver error: invalid argument

[znavidi]

Thank you for your great work!
I have been trying to fine-tune the 2B model on the toy dataset just to run the fine-tuning script properly and then adapt that to our project. I keep getting CUDA driver error: invalid argument error during runtime and even though I checked to have compatible torch (2.5.1) and cuda (tried 11.8, 12.1, and 12.4), it keeps getting this error.

python 3.12.0

torch version 2.5.1+cu121

flash-attn version 2.6.3

I appreciate any hint on where the issue might come from and how I can fix it. Thanks.