如何将我的Pytorch模块(具有就地操作)更改为可区分的？

import torch
import torch.nn as nn

class MyModule(nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()
        self.weight = nn.Parameter(torch.randn(5, 5))

    def forward(self, x):
        # 替换就地操作为原位操作
        out = torch.matmul(self.weight, x)
        return out

    def backward(self, grad_output):
        # 自定义反向传播函数
        grad_input = torch.matmul(self.weight.t(), grad_output)
        return grad_input

class MyFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight):
        # 保存输入和权重，用于反向传播
        ctx.save_for_backward(input, weight)
        output = torch.matmul(weight, input)
        return output

    @staticmethod
    def backward(ctx, grad_output):
        # 获取保存的输入和权重
        input, weight = ctx.saved_tensors
        grad_input = torch.matmul(weight.t(), grad_output)
        grad_weight = torch.matmul(grad_output, input.t())
        return grad_input, grad_weight

# 使用可区分的操作
class MyModuleDifferentiable(nn.Module):
    def __init__(self):
        super(MyModuleDifferentiable, self).__init__()
        self.weight = nn.Parameter(torch.randn(5, 5))

    def forward(self, x):
        out = MyFunction.apply(x, self.weight)
        return out

# 创建模块实例
module = MyModule()
module_diff = MyModuleDifferentiable()

# 运行前向传播
x = torch.randn(5, 5)
output = module(x)
output_diff = module_diff(x)

# 打印结果
print(output)
print(output_diff)

# 运行反向传播
grad_output = torch.randn(5, 5)
grad_input = module.backward(grad_output)
grad_input_diff, grad_weight_diff = torch.autograd.grad(output_diff, (x, module_diff.weight), grad_output)

# 打印梯度
print(grad_input)
print(grad_input_diff)
print(grad_weight_diff)