save the last checkpoint in the paligemma colab (#105)

Author: akolesnikoff

big_vision/configs/proj/paligemma/finetune_paligemma.ipynb

@@ -1121,14 +1121,31 @@
         "display(HTML(html_out))\n"
       ]
     },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "Ai0NMbAwsr0j"
+      },
+      "source": [
+        "# Save the final checkpoint"
+      ]
+    },
     {
       "cell_type": "code",
-      "source": [],
+      "execution_count": null,
       "metadata": {
         "id": "5H_3CV33_JkV"
       },
-      "execution_count": null,
-      "outputs": []
+      "outputs": [],
+      "source": [
+        "def npsave(pytree, path):\n",
+        "  names_and_vals, _ = big_vision.utils.tree_flatten_with_names(pytree)\n",
+        "  with open(path, \"wb\") as f:\n",
+        "    np.savez(f, **{k:v for k, v in names_and_vals})\n",
+        "\n",
+        "# Takes around 4 minutes\n",
+        "npsave(params, 'my-custom-paligemma-ckpt.npz')"
+      ]
     }
   ],
   "metadata": {