Fix example of restoring checkpoint in use_checkpointing.ipynb

Open sirotenko opened this issue 1 year ago • 3 comments

What does this PR do?

This is a minor issue. In the example of how to restore a partial checkpoint when calling flax.serialization.from_state_dict we better to assign the return value to the custom_target. Otherwise the state is not updated.

Checklist

[x] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you can dismiss the other checks if that's the case).
[ ] This change is discussed in a Github issue/ discussion (please add a link).
[ ] The documentation and docstrings adhere to the documentation guidelines.
[ ] This change includes necessary high-coverage tests. (No quality testing = no merge!)

Dec 18 '23 04:12 sirotenko

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.