pseudocode.py

w = mapping_network(z) # (batch, depth, dim)

nonlinear = nn.Sequential(
  nn.Linear(dim, dim, bias=False),
  nn.LeakyRelu(),
  nn.Linear(dim, init_channel_dim * 4)
)

init_image_block = nonlinear(mean(w, dim=1)).reshape(batch, 2, 2, init_channel_dim)

# i then modified my generator network to start at 2x2 instead of 4x4, but i guess you can upsample with a convtranspose2d here for simplicity
# for some reason, if i directly go from nonlinear -> 4x4, it doesn't learn