Same input and ouput 3rd dimmension in ConvTranspose3d

I’m using Transposed Convolutions in my network and I want to keep the shape of the output 3rd(indexing from 1) dimension same as in input. How can I do that with ConvTranspose3d() ?

import torch

import torch.nn as nn

x = torch.randn(1,3,1,128,128).cuda()

layer = nn.ConvTranspose3d(in_channels=3,out_channels=3,kernel_size=4,stride=2,padding=1).cuda()

output = layer(x)

print(output.shape)
Actual Output: (1,3,2,256,256)
Expected Output : (1,3,1,256,256)

You can specify the kernel_size, stride, dilation, and padding for each dimension separately and using a single scalar value is a convenient way to set all dimensions to the same value.
This should thus work:

x = torch.randn(1,3,1,128,128).cuda()

layer = nn.ConvTranspose3d(
    in_channels=3,
    out_channels=3,
    kernel_size=(1, 4, 4),
    stride=(1, 2, 2),
    padding=(0, 1, 1)).cuda()

output = layer(x)

print(output.shape)
> torch.Size([1, 3, 1, 256, 256])