Possible bug in StyleGAN training code

[wilson1yan]

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.

flaxmodels/training/stylegan2/training.py

Line 239 in edc6a85

H, W = batch['image'].shape[1], batch['image'].shape[2]

[matthias-wright]

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