Here is the model architecture i am using
class Model(torch.nn.Module):
def __init__(self):
    super(Model, self).__init__()
    
    self.conv1 = torch.nn.Conv2d(3, 30, kernel_size=3, stride=1, padding=1)
    self.relu1 = nn.ReLU()
    self.pool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
    
    self.conv2 = torch.nn.Conv2d(30, 16, kernel_size=3, stride=1, padding=1)
    self.relu2 = nn.ReLU()

    self.conv3 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu3 = nn.ReLU()
                                 
    self.conv4 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu4 = nn.ReLU()
    self.pool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0)
                                 
    self.conv5 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu5 = nn.ReLU()
    self.conv6 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu6 = nn.ReLU()
    self.conv7 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu7 = nn.ReLU()
    self.conv8 = torch.nn.Conv2d(16, 16, kernel_size=3, stride=1, padding=1)
    self.relu8 = nn.ReLU()
    self.drop2D = nn.Dropout2d(p=0.5, inplace=False)
    self.fc1 = nn.Linear(56, 2)
    
def forward(self, input):
    
    output = self.conv1(input)
    output = self.relu1(output)
    output = self.pool1(output)
    print(output.shape)
    
    
    output = self.conv2(output)
    output = self.relu2(output)
    print(output.shape)
    
    output = self.conv3(output)
    output = self.relu3(output)
    print(output.shape)
    
    output = self.conv4(output)
    output = self.relu4(output)
    output = self.pool2(output)
    
    print(output.shape)
    
    output = self.conv5(output)
    output = self.relu5(output)
    print(output.shape)
    
    output = self.conv6(output)
    output = self.relu6(output)
    print(output.shape)
    
    output = self.conv7(output)
    output = self.relu7(output)
    
    print(output.shape)
    
    output = self.conv8(output)
    output = self.relu8(output)
    
    print(output.shape)
    
    output = self.drop2D(output)
    print(output.shape)
    
    output = self.fc1(output)
    
    print(output.shape)
    return F.log_softmax(output)

model = Model()