Saving checkpoints

[raresdolga]

Hi, I am trying to save checkpoints using the following code:

options = ocp.CheckpointManagerOptions(
            max_to_keep=self.max_checkpoints,
            create=True,
            best_fn=best_loss,
            best_mode="min",
        )

        self.checkpoint_manager = ocp.CheckpointManager(
            os.path.join(self._out_dir, "checkpoints"),
            options=options,
            item_names=("state", "metadata"),
            item_handlers={
                "state": ocp.StandardCheckpointHandler(),
                "metadata": ocp.JsonCheckpointHandler(),
            },
        )

The problem happens when I try to create the folder. It checks for multiprocessing using a flag, but the flag is not parsed. Error is:

absl.flags._exceptions.UnparsedFlagAccessError: Trying to access flag --experimental_orbax_use_distributed_process_id before flags were parsed.

Seems to be triggered by this line of code in the library:

absl.flags._exceptions.UnparsedFlagAccessError: Trying to access flag --experimental_orbax_use_distributed_process_id before flags were parsed.

I should note that I am running my code from a child process.

Do you have any suggestions on how to avoid this? I am not sure why the flag is not picked up since it has a default value. I tried parsing them from the main file using args, but then it asks to define my own arguments as flags, which seems as an overhead.

[cpgaffney1]

This problem should have been fixed by now, let me know if you are still seeing it.