import torch
from torch.utils import checkpoint

def fun(x, theta, t):
    return theta['alpha']**t + theta['beta']**t

def wrapper(x):
    x = torch.tensor(x, requires_grad=True)
    theta = {}
    theta['alpha'] = x[0]
    theta['beta'] = x[1]
    out = torch.empty(3)
    for t in range(3):
        out[t] = checkpoint.checkpoint(fun, *(x, theta, t))
    y = out.sum()
    y.backward()
    grad = x.grad
    return (y.item(), grad)

print(wrapper([2., 3.]))

If I modify  fun to  also take x, I still get an a type error saying “save_for_backward can only save variables, but argument 1 is of type dict.”  Does this mean that the function in a checkpoint can only take tensors for which requires_grad=True? If so that seems extremely limiting. I don’t even know how this simple example can be rewritten to do that.