Update infer.py. Let process 0 write the per shard 'COMPLETED' indicator file.
Update infer.py. Let process 0 write the per shard 'COMPLETED' indicator file.

PiperOrigin-RevId: 465619923

google-research/t5x