Tutorial 9 (JAX): Fixing jitting of bind

Author: phlippe

docs/tutorial_notebooks/JAX/tutorial9/AE_CIFAR10.ipynb

@@ -6649,7 +6649,7 @@
     "    \n",
     "    @jax.jit\n",
     "    def encode(imgs):\n",
-    "        return trainer.model_bd.encoder(imgs)\n",
+    "        return trainer.model.bind({'params': trainer.state.params}).encoder(imgs)\n",
     "    \n",
     "    for imgs, _ in tqdm(data_loader, desc=\"Encoding images\", leave=False):\n",
     "        z = encode(imgs)\n",