import torch


def diff_single(y, x, device):
    """Computes the derivative of a single y w.r.t. x
    """

    yp = torch.autograd.grad(
        y,
        x,
        grad_outputs=torch.ones_like(y, device=device),
        create_graph=True,
        retain_graph=True,
    )[0]
    return yp.transpose(0, 1)