# Codeblock 2
class ConvBlock(nn.Module):
    def __init__(self, 
                 in_channels, 
                 out_channels, 
                 kernel_size, 
                 stride, 
                 padding, 
                 maxpool_flag=False):
        super().__init__()
        self.maxpool_flag = maxpool_flag
        
        self.conv = nn.Conv2d(in_channels=in_channels,       #(1)
                              out_channels=out_channels, 
                              kernel_size=kernel_size, 
                              stride=stride, 
                              padding=padding)
        self.leaky_relu = nn.LeakyReLU(negative_slope=0.1)   #(2)
        
        if self.maxpool_flag:
            self.maxpool = nn.MaxPool2d(kernel_size=2,       #(3)
                                        stride=2)
            
    def forward(self, x):
        print(f'original\t: {x.size()}')

        x = self.conv(x)
        print(f'after conv\t: {x.size()}')
        
        x = self.leaky_relu(x)
        print(f'after leaky relu: {x.size()}')
        
        if self.maxpool_flag:
            x = self.maxpool(x)
            print(f'after maxpool\t: {x.size()}')
        
        return x