I’m not sure how linearZ is supposed to work, as it seems like you would like to calculate the gradients in backward as well as manipulate some model parameters.
Could you explain the use case a bit as I think both steps (gradient calculation and parameter manipulation) should be done separately.
In case you would just like to reimplement a simple linear function, this code should work:
class linearZ(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, weight):
        ctx.save_for_backward(input, weight)
        output = input.mm(weight.t())
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, weight = ctx.saved_tensors
        grad_input = grad_output.mm(weight)
        grad_weight = grad_output.t().mm(input)
        return grad_input, grad_weight


class MyLinearZ(nn.Module):
    def __init__(self):
        super(MyLinearZ, self).__init__()
        self.fn = linearZ.apply
        self.weight = nn.Parameter(torch.randn(1, 1) * 5.66)

    def forward(self, x):
        x = self.fn(x, self.weight)
        return x