Restart fix

Get rid of the error when running with checkpoints + MPI.