如何在nn.Sequential模型中使用自定义torch.autograd.Function

import torch
import torch.nn as nn
import torch.autograd as autograd

# Step 1: 定义自定义的torch.autograd.Function类
class MyFunction(autograd.Function):
    @staticmethod
    def forward(ctx, input):
        # 前向传播的计算逻辑
        ctx.save_for_backward(input)
        output = input.clamp(min=0)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # 反向传播的计算逻辑
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input

# Step 2: 创建包含自定义Function的模块
class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()

    def forward(self, input):
        return MyFunction.apply(input)

# Step 3: 在nn.Sequential模型中使用自定义模块
model = nn.Sequential(
    MyModule(),
    nn.Linear(10, 5),
    nn.ReLU(),
    nn.Linear(5, 2)
)

# 使用模型进行前向传播
input = torch.randn(1, 10)
output = model(input)
print(output)