ptrblck:

Could you post a code snippet


Sure, the class is the following:
class Generator(nn.Module):
    def __init__(self,n_gpu):
        super(Generator,self).__init__()
        self.n_gpu = n_gpu
        self.G_ch_mul = G_ch_mul  #Channel width multiplier
        self.resolution = resolution    # Resolution of the OUTPUT
        self.G_depth = G_depth          # Number of resblocks per stage
        self.arch = G_arch(self.G_ch_mul)[resolution]
        self.G_init = G_init

        self.linear = nn.Linear(noise_dim,self.arch['in_channels'][0]*(4**2))  #(N,noise-dim) > (N,64*16*4*4)

        self.blocks = []
        for index in range(len(self.arch['out_channels'])):                  #In total we have 5 cycles
          self.blocks += [[GBlock(input_ch=self.arch['in_channels'][index],  #(16,16,8,4,2) * 64
                                  output_ch=self.arch['in_channels'][index] if g_index==0 else self.arch['out_channels'][index], #((16,16),(16,8),(8,4),(4,2),(2,1))*64  (O1,O2)
                                  upsample=(functools.partial(F.interpolate, scale_factor=2) if self.arch['upsample'][index] and g_index == (self.G_depth-1) else None))]
                         for g_index in range(self.G_depth)]

          if self.arch['attention'][self.arch['resolution'][index]]:
            print('Adding attention layer in G at resolution %d' % self.arch['resolution'][index])
            self.blocks[-1] += [Attention(self.arch['out_channels'][index], SNConv2d)]

        self.blocks = nn.ModuleList([nn.ModuleList(block) for block in self.blocks])

        #self.output_layer = nn.Sequential(nn.BatchNorm2d(self.arch['out_channels'][-1]),
        #                                  nn.ReLU(True),
        #                                  nn.Conv2d(self.arch['out_channels'][-1],channels,kernel_size=3,padding=1))

        self.bn = nn.BatchNorm2d(self.arch['out_channels'][-1])
        self.act = nn.ReLU(True)
        self.l_conv = nn.Conv2d(in_channels=self.arch['out_channels'][-1],out_channels=channels,kernel_size=3,padding=1)
     
    def forward(self,z,y):

        h = self.linear(z)   #(N,noise_dim) > (N,64*4*4*16)
        h = h.view(h.size(0),-1,4,4)  #(N,64*4*4*16) > (N,16*64,4,4)
        for index, blocklist in enumerate(self.blocks):
            # Second inner loop in case block has multiple layers
            for block in blocklist:
              h = block(h, y)
        out = self.bn(h)
        out = self.act(h)
        out = self.l_conv(h)
        out = torch.tanh(out)

        return out

I don’t know how useful it is since the same line works well when it is outside my trianing loop, the in_channels=128, out_channels=1, Win=128 and Hin=128