Add flax text class colab (#11824)

Commit

4 years ago

Add flax text class colab (#11824) * fix_torch_device_generate_test * remove @ * add flax glue link