All layers translators with nan loss while training the lens

[AmeenAli]

Hello!

Thanks for sharing this amazing work!

I am trying to train the lens over a new dataset HF Dataset (note that the original "the pile" dataset was removed from the internet because of a DMCA claim - as i understood it. as follows :

tuned-lens train --model.name EleutherAI/pythia-160m-deduped --output ./output/ --per_gpu_batch_size=6 --data.name HuggingFaceH4/ultrachat_200k --split train_gen --text_column prompt --wandb test

However, once I look into wandb logs i notice that loss/translator_1 - 12 are all NaNs also the weight_norms_1-12. I have also tried different datasets, but the issue still exists. Any idea what is off?

Thank you!