# Codeblock 6
class FirstTransition(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        
        self.bn   = nn.BatchNorm2d(num_features=in_channels)
        self.relu = nn.ReLU()
        self.conv = nn.Conv2d(in_channels=in_channels, 
                              out_channels=out_channels,   #(1)
                              kernel_size=1,               #(2)
                              padding=0,
                              bias=False)
        self.dropout = nn.Dropout(p=0.2)
     
    def forward(self, x):
        print(f'original\t\t: {x.size()}')
        
        out = self.dropout(self.conv(self.relu(self.bn(x))))
        print(f'after first_transition\t: {out.size()}')
        
        return out