Ok, let’s talk about a simpler example:
import torch
from torch.utils import checkpoint

def fun(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, *(theta, t))
    y = out.sum()
    y.backward()
    grad = x.grad
    return (y.item(), grad)

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

From your response, it sounds like maybe I should make a checkpoint at each iteration of my for loop in logL, so here is a toy example where I am essentially trying to do the same thing in the function wrapper. When I run this code, it returns “RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn”, but theta requires grad here right? What am I doing wrong? Note that if  checkpoint.checkpoint() is removed, the function returns the appropriate gradient.