Hi!
I need to create a 2-layer convolutional net that takes as input a 3-224-224 image, uses 50 kernels of 33,  50 kernels of 44 and 50 kernels of 5*5 in each layer to perform convolutions and then returns an image.
Can someone help me with an example of this? I tried the following code with a batch size of 16 but the output is torch.Size([16, 142572]), not the 16-3-224-224 that I was expecting.
 class Net(nn.Module):
     def __init__(self):
         super(PrepNetwork, self).__init__()
         self.layer1P = nn.Sequential(
             nn.Conv2d(3, 50, kernel_size=3, padding=1),
             nn.ReLU(),
             nn.Conv2d(50, 50, kernel_size=4, padding=1),
             nn.ReLU(),
             nn.Conv2d(50, 50, kernel_size=5, padding=1),
             nn.ReLU())
         self.layer2P = nn.Sequential(
             nn.Conv2d(50, 50, kernel_size=3, padding=1),
             nn.ReLU(),
             nn.Conv2d(50, 50, kernel_size=4, padding=1),
             nn.ReLU(),
             nn.Conv2d(50, 3, kernel_size=5, padding=1))
 
     def forward(self, x):
         h1 = self.layer1P(x)
         out = self.layer2P(h1)
         out = out.view(out.size(0), -1)
         return out

Thanks in advance!