import torch
from torch.autograd import Function

class VarianceMatch(Function):
    @staticmethod
    def forward(ctx, x):
        y : torch.Tensor = x * 2
        y.requires_grad_()
        y.retain_grad() # Does not work.
        def h(grad): # y.grad is None so this gives me an exception.
            return grad * torch.sqrt(torch.square(y.grad).sum(1) / torch.square(grad).sum(1))
        x.register_hook(h)
        return y

    @staticmethod
    def backward(ctx, grad_y):
        return grad_y * 2

i = torch.scalar_tensor(2,requires_grad=True)
x = VarianceMatch.apply(i)
x.backward(torch.scalar_tensor(1))
print(i.grad)

I am trying to hack it using hooks, but that is not working for me either since the output y does not have the gradient. At this point I am out of ideas.