Should be fixed now, i just need to update the nb but the Accelerate specific bugs should be solved outside of that specific cause.
Declare the model outside the training function, and pass it in as an argument instead. A PR with this fix will go live tommorow as well.
The model should only be declared once on a TPU and its passed back and forth, rather than duplicated in a notebook/forked process