Ho right. because you “hide” these inputs from the checkpoint code, it cannot properly detach them during the backward…
This is code that I haven’t run so there might be typos in it 
def detach(obj):
    # A more general detach that opens up dict and tuples
    if torch.is_tensor(obj):
        new = obj.detach()
        new.requires_grad_(obj.requires_grad)
        return new
    elif isinstance(obj, tuple):
        return tuple(detach(el) for el in obj)
    elif isinstance(obj, dict):
        return {k:detach(v) for k, v in obj.items()}
    else:
        return obj


    def fn(x, data=data, t=t):
        data = detach(data)
        t = detach(t)
        model.update(data, t)
        return model.p_a

That will solve this issue.
This checkpoint wasn’t really made for anything but neural nets.Sorry it is so painful