So we are using a residual block for our model, the block looks like this:
class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, triplet=True):
        super(ResidualBlock, self).__init__()
        
        # conv layers
        self.conv_res = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)
        self.conv_1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.conv_2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        
        # initialization
        kaiming_normal_(self.conv_res.weight, nonlinearity='relu')
        kaiming_normal_(self.conv_1.weight, nonlinearity='relu')
        kaiming_normal_(self.conv_2.weight, nonlinearity='relu')
        
        # batch norm layers
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.bn3 = nn.BatchNorm2d(out_channels)
        
        # activations 
        self.relu = nn.ReLU()

        # flags
        self.triplet = triplet

    def forward(self, x):
        
        # out placeholder
        out = None
        
        # if the block is a triplet
        if self.triplet:
            
            # first stage
            out_1 = self.conv_1(x)
            out_1 = self.bn1(out_1)
            out_1 = self.relu(out_1)
            
            # second stage
            out_2 = self.conv_2(out_1)
            out_2 = self.bn2(out_2)
            
            # third stage - residual connection
            out_res = self.conv_res(x)

        # if the block is a twin
        else:
            
            # first stage
            out_1 = self.conv_1(x)
            out_1 = self.bn1(out_1)
            out_1 = self.relu(out_1)
            
            # second stage
            out_2 = self.conv_2(out_1)
            out_2 = self.bn2(out_2)
            
        # add the activations
        if self.triplet:
            out = out_2 + out_res
        else:
            out = out_2 + x

        # final activation function
        out = self.relu(out)
            
        return out

The model is basically a stack of this residual blocks. Here is the encoder part of it:
class Encoder(nn.Module):
    def __init__(self):
        super(Encoder, self).__init__()
        
        # conv layer
        # using 3 by 3 receptive field
        self.conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, stride=1, padding=1)
        
        # initialize the conv layer
        kaiming_normal_(self.conv.weight.data, nonlinearity='relu')  # init with kaiming he

        # residual blocks
        self.res1 = ResidualBlock(in_channels=16, out_channels=32)
        self.res2 = ResidualBlock(in_channels=32, out_channels=64)
        self.res3 = ResidualBlock(in_channels=64, out_channels=128)
        self.res4 = ResidualBlock(in_channels=128, out_channels=256)
        self.res5 = ResidualBlock(in_channels=256, out_channels=512)

        # fully connected layer
        self.fc1024 = nn.Linear(in_features=8192, out_features=1024)
        
        # initialize the fc
        kaiming_normal_(self.fc1024.weight.data, nonlinearity='relu')
        
        # pool layer
        self.pool = nn.AvgPool2d(kernel_size=2)
        
    def forward(self, x):
        
        # x is 3x256x256
        
        # input block
        out = self.conv(x)    # 16x256x256
        out = self.pool(out)  # 16x128x128
        
        print("CONV MAX ACT: ", out.max())
        
        # residual block 1
        out = self.res1(out)  # 32x128x128
        out = self.pool(out)  # 32x64x64
        
        print("RES1 MAX ACT: ", out.max())
        
        # residual block 2
        out = self.res2(out)  # 64x64x64
        out = self.pool(out)  # 64x32x32
        
        print("RES2 MAX ACT: ", out.max())
        
        # residual block 3
        out = self.res3(out)  # 128x32x32
        out = self.pool(out)  # 128x16x16
        
        print("RES3 MAX ACT: ", out.max())
        
        # residual block 4
        out = self.res4(out)  # 256x16x16
        out = self.pool(out)  # 256x8x8
        
        print("RES4 MAX ACT: ", out.max())
        
        # residual block 5
        out = self.res5(out)  # 512x8x8
        out = self.pool(out)  # 512x4x4
        
        print("RES5 MAX ACT: ", out.max())
        
        print("CONV MAX WEIGHT: ", self.conv.weight.max(), self.conv.bias.max())
        print("RES1 MAX WEIGHTS: ", self.res1.conv_1.weight.max(), self.res1.conv_2.weight.max(), self.res1.conv_res.weight.max())
        print("RES2 MAX WEIGHTS: ", self.res2.conv_1.weight.max(), self.res2.conv_2.weight.max(), self.res2.conv_res.weight.max())
        print("RES3 MAX WEIGHTS: ", self.res3.conv_1.weight.max(), self.res3.conv_2.weight.max(), self.res3.conv_res.weight.max())
        print("RES4 MAX WEIGHTS: ", self.res4.conv_1.weight.max(), self.res4.conv_2.weight.max(), self.res4.conv_res.weight.max())
        print("RES5 MAX WEIGHTS: ", self.res5.conv_1.weight.max(), self.res5.conv_2.weight.max(), self.res5.conv_res.weight.max())
        
        
        print("OUT_PRE MAX: ", out.max())
        
        # reshape
        out = out.view((out.shape[0], -1))

        # fully connected
        out = self.fc1024(out)
        
        print("OUT_SHAPE: ", out.shape, "OUT: ", out, "OUT MAX: ", out.max())
        
        # get the mean and log-variance
        mu = out[:, :512]
        logvar = out[:, 512:]
        
        # reparametrization trick
        z = self.reparameterize(mu, logvar)
        
        # return the latent vector, mu and logvar
        return z, mu, logvar

The print statements print the next stats:

image.png2017×800 156 KB

This is on the first batch of the first epoch.