Distillation code works on TPU?

Wondering if anyone has tested the distillation code for Seq2Seq models on a TPU?
Just want to know if there has been a positive/negative experience on this.

@sshleifer you probably know more about this.

Thanks!