Loading checkpoints when models built using a 'setup' block

4

Hey @jirka, thanks for taking a look. I fleshed out the example above in this Colab:

Please let me know if everything makes sense here. Note, cell [5] will error, as expected and shows the recommended load method fails.