Fix UNet segmentation example

[xrhd]

Resolves #113

Reference

• jax.extend.backend.get_backend
• Orbax checkpoints

Checklist

• I have read the Contribution Guidelines and used pre-commit hooks to format this commit.
• [] I have added all the necessary unit tests for my change. (run_model.py for model usage, test_outputs.py and/or model_validation_colab.ipynb for quality).
• [] (If using an LLM) I have carefully reviewed and removed all superfluous comments or unneeded, commented-out code. Only necessary and functional code remains.
• I have signed the Contributor License Agreement (CLA).