Your new code runs fine but it’s still returning
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.
So if I understand you correctly, the problem is that checkpoint isn’t detaching all of the tensors for which requires_grad=True.  Could it be that the tensor attributes of  model need to be detached as well?