请使用pytorch框架编写深度可分离卷积代码

以下是使用PyTorch框架编写深度可分离卷积的示例代码：

import torch
import torch.nn as nn

class SeparableConv2d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size=3, stride=1, padding=1):
        super(SeparableConv2d, self).__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=stride, padding=padding, groups=in_channels)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0)
        
    def forward(self, x):
        x = self.depthwise(x)
        x = self.pointwise(x)
        return x

在这个示例中，我们定义了一个名为`SeparableConv2d`的类，它继承自PyTorch的`nn.Module`类。该类实现了深度可分离卷积，包含两个卷积层：深度卷积层和点卷积层。深度卷积层实现了深度可分离卷积的深度卷积部分，通过设置`groups=in_channels`，我们将输入通道分组为一个个单独的通道，然后对每个通道进行卷积操作。点卷积层实现了深度可分离卷积的点卷积部分，它将所有通道的卷积结果合并起来，生成最终的输出结果。

在`forward`函数中，我们首先对输入进行深度卷积操作，然后将卷积结果输入到点卷积层中，最终输出卷积结果。