Fix checkpointing bug