Latent Reasoning VRAM Constrained model

I had to squeeze every mb i could and i managed to get the model seemingly progressing, tho eventually i've hit OOM and i decided to give up.

I'll start a branch where i can train this on TPUs on Google Cloud (in small runs to prove the model works)

If y'all could evaluate my code that'd be awesome