Hey there!
I’ve been playing around with plotting images generated from StyleGAN3, but there seems to be something funky going on with the color channels that i can’t really pin down.
Generating an image with:
device = torch.device('cuda')
# open pickle file and load generator
network_pkl = './models/stylegan3-r-afhqv2-512x512.pkl'
with dnnlib.util.open_url(network_pkl) as f:
    G = legacy.load_network_pkl(f, force_fp16=True)['G_ema'].cuda()

noise_mode = 'const'
trunc = 1
seed = 10
z = torch.from_numpy(np.random.RandomState(seed).randn(1, G.z_dim)).to(device)
# returns [b, c, 512, 512] tensor
img = G(z, truncation_psi=trunc, noise_mode=noise_mode, c=None)
# permute to [b, 512, 512, c] and scale to 0-255
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
# remove batch channel
img = img.squeeze(0) 
# move tensors to cpu
img_cpu = img.cpu()
# plot with pyplot
plt.imshow(img_cpu.numpy())
# save with pillow
PIL.Image.fromarray(img_cpu.numpy(), 'RGB').save(f'seed{seed}.png')

From my understanding the generator of StyleGAN3 returns an image as a tensor of [batch_size, channels, height, width] with channels being RGB. The plotted image seems to be have completely wrong color information.
From plt.imshow(img_cpu.numpy()):

Saving with PIL.Image.fromarray(img_cpu.numpy(), 'RGB').save(f'seed{seed}.png'):

Screenshot from 2021-10-18 15-34-32516×519 87.7 KB

I swapped around the channels to BGR, GRB etc. but couldn’t get the plot to display the correct image. Would appreciate any help with this :).
Cheers
Edit: Issue wasn’t code related. Was using jetbrains DataSpell ide which is in early preview and still has some bugs.