Hey @ tmabraham ! This project sounds interesting! I would love to join it. Although I have not worked with JAX yet, I have experience with GANs and Computer Vision using PyTorch. It would be a great learning experience for me. Please let me know how can I join?