* changes to allow for tpu training * black * tpu * tpu
* change masking to direct labelings * fix black * switch to ignore index * . * fix black