mlterpstra-umc:

I thought this could be done with torch.nn.functional.conv2d(images, kernel.expand(2048, 1, 2, 2)) but this results a final matrix of size [2048, 2048, 11, 11].


Yes, that’s right. I did not notice that the output size should be [2048, 1, 12, 12], so specifically 1 output channel. So the convolution kernel shape is [output_channels, input_channels, kernel_dim_0, kernel_dim_1].
So determining the kernel size:

Desired output channels: 1

Input channels: 1

Filter size: 2x2


Therefore,  the kernel size for this should just be [1, 1, 2, 2]:
>>> a = torch.randn(2048, 1, 12, 12)
>>> pad = nn.ZeroPad2d((1, 0, 1, 0))
>>> ap = pad(a)
>>> ap.shape
torch.Size([2048, 1, 13, 13])
>>> kernel = torch.randn(1, 1, 2, 2)
>>> kernel.shape
torch.Size([1, 1, 2, 2])
>>> b = F.conv2d(ap, kernel)
>>> b.shape
torch.Size([2048, 1, 12, 12])