search

matthias-wright / flaxmodels

Pretrained deep learning models for Jax/Flax: StyleGAN2, GPT2, VGG, ResNet, etc.
https://github.com/matthias-wright/flaxmodels
237 stars 25 forks source link

Possible bug in StyleGAN training code #31

Open wilson1yan opened 1 year ago

wilson1yan commented 1 year ago

Hi,

Nice repo! Is this line a bug? Since I think batch['images'] is N x B x H x W x C, so the indices should be shift up by 1.

https://github.com/matthias-wright/flaxmodels/blob/edc6a8571a6d7202bd9f3bc9241221405c083fd4/training/stylegan2/training.py#L239

matthias-wright commented 1 year ago

Hi @wilson1yan, sorry for the late reply! I think you are right, do you want to submit a PR?