As a workaround I get the Jacobian for each index of Y. This is computationally intensive. But does the work for now.
def squared_grad(self, Y, X, G):
        result = torch.zeros_like(X)
        g = torch.zeros_like(G)
        g_np = g.data.cpu().detach().numpy()
        gout = torch.ones(1).cuda()[0]
        for index, e in np.ndenumerate(g_np):
            result += torch.autograd.grad(Y[index], X, grad_outputs=gout, retain_graph=True)[0]**2*G[index]
        return result

Any comments?
Thanks.