access state from restored composite args after restore

[keshavb96]

manager.restore now returns a orbax.checkpoint._src.handlers.composite_checkpoint_handler object instead of a PyTree, so we need to access the state member of the returned object.