import torch


def grad(x, y, model, params, criterion):
    """Compute new parameters with computation graph intact.

     Arguments:
        x (torch.Tensor): input tensor.
        y (torch.Tensor): target tensor.
        model (Warp): warped model.
        params (list): list of parameters to differentiate
            the loss with respect to.
        criterion (fun): task loss criterion.
    """
    device = next(model.parameters()).device

    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)

    p = model(x, cache_parameters=False)
    loss = criterion(p, y)
    return torch.autograd.grad(loss, params, create_graph=True)