Nx While loop training slow down when passing through frozen embedings

[jbc5099]

I’m using EXLA version 0.11.0, and have rocm set as backend and EXLA set as the compiler. Here’s a gist I made from the .livemd file: https://gist.github.com/Brady-Conn/810f8a9325526b6ecae2b86fbaf7e393.js