Small update to checkpoint fix.

The previous commit was still incomplete. This commit hopefully
fixes this for good. Code passes the testsuite.