Conversion of (16,7,2,16,768) to (16,7,2,768) using CNN

Hi, i am trying to Convert tensor of (16,7,2,16,768) to (16,7,2,768) using CNN. How can i apply CNN on only last two layers so that the value of the other dimensions are untouched.

Thank you.

Could you explain a bit what “untouched” would mean in this context?
If you are concerned about the size of the other dimensions, you could use e.g. nn.Conv3d with a kernel size of 16 for the second dimension of the kernel and squeeze this dimension afterwards:

x = torch.randn(16, 7, 2, 16, 768)
conv = nn.Conv3d(7, 7, kernel_size=(1, 16, 1))
out = conv(x)
print(out.shape)
> torch.Size([16, 7, 2, 1, 768])
out = out.squeeze(3)
print(out.shape)
> torch.Size([16, 7, 2, 768])

However, this would of course use all provided values of the input tensor.