Now I’m getting the following error:
RuntimeError: Trying to backward through the graph a second time, but the buffers have already been freed. Specify retain_graph=True when calling backward the first time.
I only call .backward() once, and setting retain_graph=True just returns the same error. Here is my modified code:

def logL(x, theta, model, data):
    N, S, T = (data['N'], data['S'], data['T'])
    p_a = torch.ones(N * S)
    model.start(theta, data['mat'])
    for t in range(1, T + 1):
        def fn(x, data=data, t=t):
            model.update(data, t)
            return model.p_a
        out = checkpoint.checkpoint(fn, x)
        p_a = p_a * (npeat(data['obs'][:,:,t-1], (S, 1)) * out).sum(1)   
    ll = p_a.reshape(N, S).mean(1).log().sum()   
    return ll 

def fun(x, cons, est, trans, model, data):
    print(x)
    for con in cons:
        valid = np.all(con['fun'](x, *con['args']) > 0)
        if valid == False: break
    if valid == False:
        ll = float('nan')
        grad = np.empty(len(x))
        grad[:] = np.nan
        print('Constraint violation')
        print()
        return (-ll, -grad)     
    else:
        torch.manual_seed(7)
        ex0 = (None,)
        ex1 = (slice(None), None)
        x = torch.tensor(x.astype(np.float32), requires_grad=True)
        ll = 0
        for d in range(len(data)):
            N, S, K = (data[d]['N'], data[d]['S'], data[d]['K'])
            p = 0
            mu = torch.tensor([])
            re = []
            theta = {}
            for (key, val) in est.items():
                if type(val) != str: 
                    theta[key] = torch.tensor(val)
                elif val == 'pool':
                    theta[key] = x[p][ex0]
                    p += 1
                elif val == 'rand':
                    mu = torch.cat((mu, x[p][ex0]))
                    p += 1
                    re.append(key)
            R = len(re)
            di = tuple([torch.tensor(ind, dtype=torch.long) 
                        for ind in np.diag_indices(R)])
            tril = tuple([torch.tensor(ind, dtype=torch.long) 
                          for ind in np.tril_indices(R, -1)])
            triu = tuple([torch.tensor(ind, dtype=torch.long)
                          for ind in np.triu_indices(R, 1)]) 
            sigma = torch.empty(R, R)
            sigma[di] = x[p:p+R]
            sigma[tril] = x[p+R:]
            sigma[triu] = sigma.t()[triu]                
            sample = MultivariateNormal(mu, sigma).rsample((N * S,))
            theta.update({re[r]: trans[re[r]](sample[:,r]) for r in range(R)})
            theta['upsilon'] = torch.zeros(N * S)[ex1]
            if K > 1: 
                theta['upsilon'] = torch.cat((theta['upsilon'], 
                                              theta['upsilon2'][ex1]), 1)
            if K > 2: 
                theta['upsilon'] = torch.cat((theta['upsilon'], 
                                              theta['upsilon3'][ex1]), 1)
            theta['upsilon'] = theta['upsilon'] / theta['kappa']
            if data[d]['treat'] == 'pvd': theta['omega'] = torch.zeros(1)
            ll = ll + logL(x, theta, model, data[d])
        ll = ll / 1e+2
        ll.backward()
        grad = np.array(x.grad.cpu())
        print(ll.item() * 1e+2)
        print()
        return (-ll.item(), -grad)   

Note that fun has only been changed to pass the parameter vector x to logL, and logL has been changed to reflect the changes we made in the toy example.