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).

[chapman20j]

Hi @xrhd. Thanks for this commit. Just leaving some brief comments.

1. Could you remove emojis from the notebook
2. Could you update the intro information to reflect that this is no longer a work in progress
3. In section 4, the script re-loads the weights and then trains again. The further trained weights aren't used in the final inference outputs. Could you just use the weights from section 4 in section 6?