This should give you a start for the automatically calculated gradient bit:
class MyFn(torch.autograd.Function):
    @staticmethod
    def forward(ctx, a_):
        with torch.enable_grad():
            a = a_.detach().requires_grad_()
            res = a**2
        ctx.save_for_backward(a, res)
        return res.detach()
    @staticmethod
    def backward(ctx, grad_out):
        a, res = ctx.saved_tensors
        gr, = torch.autograd.grad(res, a, grad_out, retain_graph=True)
        return gr

x = torch.randn(2,2, requires_grad=True, dtype=torch.double)
print(torch.autograd.gradcheck(MyFn.apply, (x,)))

I must admit I have no idea why you’d need to specify “retain_graph”, it’s just the empirical observation that without it, it doesn’t run.
Best regards
Thomas