net = nn.Linear(out_, in_, bias=False)

class Function(T.autograd.Function):

    @staticmethod
    def forward(ctx, input, net):
 
        output = net(input)
        ctx.save_for_backward(net, input)
        return output

    @staticmethod
    def backward(ctx, grad_output):...

How to use a net inside autograd.Function? Is it possible? What are my options?