Flax support?

[Kimiko-AI]

Maybe worth considering since google give away TPUs almost for me with TRC program.

[ZeroCool940711]

I think we have TPU support, might not be the best way to do it but I think we do support them. I think someone was working on improving it some time ago, not sure how that went tho.

Btw, if you have access to some TPUs or other hardware and you want to donate it for training we would appreciate it, we are actually in need of better hardware to train a base model to release open source for everyone to use, we currently are training on an RTX 3050 with 8GB of VRAM which as you can guess is really slow and it limits what we can do.

[Kimiko-AI]

How much data are you going to use? I can pull some string to get some TPUs.

[ZeroCool940711]

How much data are you going to use? I can pull some string to get some TPUs.

We have two datasets that we can use for training, one is 100-200GB and the other one is 1TB or so in size, we can use the smaller one if storage is limited, we also have a smaller one we can use but the bigger the dataset the better, specially for training the VAE since it will increase the quality of the model a lot.