`pw_checkpoints` undefined in `train.py`

[yhtang]

When do training, we hit the error:

  File "/opt/t5x/t5x/train.py", line 213, in train
    checkpoint_cfg.restore.checkpointer_cls, pw_checkpoints.Checkpointer
NameError: name 'pw_checkpoints' is not defined

The code path appears to exist for a long time but has not been run by default until recently.