Implement SeparableConv2D in Pytorch

Main objective

PyTorch equivalent for SeparableConv2D with padding = 'same':

from tensorflow.keras.layers import SeparableConv2D
x = SeparableConv2D(64, (1, 16), use_bias = False, padding = 'same')(x)

What is the PyTorch equivalent for SeparableConv2D?

This source says:

If groups = nInputPlane, kernel=(K, 1), (and before is a Conv2d layer with groups=1 and kernel=(1, K)), then it is separable.

While this source says:

Its core idea is to break down a complete convolutional acid into a two-step calculation, Depthwise Convolution and Pointwise.

This is my attempt:

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, depth, kernel_size, bias=False):
        super(SeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, out_channels*depth, kernel_size=kernel_size, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(out_channels*depth, out_channels, kernel_size=1, bias=bias)

    def forward(self, x):
        out = self.depthwise(x)
        out = self.pointwise(out)
        return out

Is this correct? Is this equivalent to tensorflow.keras.layers.SeparableConv2D?

What about padding = 'same'?

How to ensure that my input and output size is the same while doing this?

My attempt:

x = F.pad(x, (8, 7, 0, 0), )

Because the kernel size is (1,16), I added left and right padding, 8 and 7 respectively. Is this the right way (and best way) to achieve padding = 'same'? How can I place this inside my SeparableConv2d class, and calculate on the fly given the input data dimension size?

All together

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, depth, kernel_size, bias=False):
        super(SeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, out_channels*depth, kernel_size=kernel_size, groups=in_channels, bias=bias)
        self.pointwise = nn.Conv2d(out_channels*depth, out_channels, kernel_size=1, bias=bias)

    def forward(self, x):
        out = self.depthwise(x)
        out = self.pointwise(out)
        return out


class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.separable_conv = SeparableConv2d(
            in_channels=32, 
            out_channels=64, 
            depth=1, 
            kernel_size=(1,16)
        )
        
    def forward(self, x):
        x = F.pad(x, (8, 7, 0, 0), )
        x = self.separable_conv(x)
        return x

Any problem with these codes?