2 ms·
What library did you use for JAX on TPU? Also curious how much data you had?
by totoglazer 5y ago
What library did you use for JAX on TPU? Also curious how much data you had?
- rg111 5y agoWe used Flax, of course. We thought Haiku was cool, but none of us knew it, and learning resources were scarce. But we are very happy with Flax. We used the Bengali subset of the mC4 dataset [0] for training the GPT-2 model. [0]: https://huggingface.co/datasets/mc4 https://huggingface.co/datasets/mc4