# Codeblock 2
class CNN(nn.Module):
    def __init__(self):
        super().__init__()
        
        self.relu = nn.ReLU()
        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
        
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(in_channels=64, out_channels=256, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(in_channels=256, out_channels=512, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(in_channels=512, out_channels=1024, kernel_size=3, padding=1)
        self.conv5 = nn.Conv2d(in_channels=1024, out_channels=2048, kernel_size=3, padding=1)
        
    def forward(self, x):
        print(f'original\t: {x.size()}\n')
        
        x = self.relu(self.conv1(x))
        print(f'after conv1\t: {x.size()}')
        
        x = self.maxpool(x)
        print(f'after pool\t: {x.size()}\n')
        
        x = self.relu(self.conv2(x))
        print(f'after conv2\t: {x.size()}')
        
        x = self.maxpool(x)
        print(f'after pool (c2)\t: {x.size()}\n')
        
        c2 = x.clone()
        
        x = self.relu(self.conv3(x))
        print(f'after conv3\t: {x.size()}')
        
        x = self.maxpool(x)
        print(f'after pool (c3)\t: {x.size()}\n')
        
        c3 = x.clone()
        
        x = self.relu(self.conv4(x))
        print(f'after conv4\t: {x.size()}')
        
        x = self.maxpool(x)
        print(f'after pool (c4)\t: {x.size()}\n')
        
        c4 = x.clone()
        
        x = self.relu(self.conv5(x))
        print(f'after conv5\t: {x.size()}')
        
        c5 = self.maxpool(x)
        print(f'after pool (c5)\t: {c5.size()}\n')
        
        return c2, c3, c4, c5