Update in notebooks

4 files changed, 296 insertions, 47 deletions

diff --git a/notebooks/01-look-at-emnist.ipynb b/notebooks/01-look-at-emnist.ipynb
index 5e5750e..f84bef2 100644
--- a/notebooks/01-look-at-emnist.ipynb
+++ b/notebooks/01-look-at-emnist.ipynb
@@ -6,45 +6,7 @@
    "metadata": {
     "tags": []
    },
-   "outputs": [
-    {
-     "ename": "ImportError",
-     "evalue": "cannot import name 'Schema' from 'pydantic' (/home/aktersnurra/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/pydantic/__init__.py)",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning/app/__init__.py:7\u001b[0m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_utilities\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mimports\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m module_available\n\u001b[0;32m----> 7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _root_logger  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _logger  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/__init__.py:31\u001b[0m\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m__version__\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m version \u001b[38;5;28;01mas\u001b[39;00m __version__  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n\u001b[0;32m---> 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningApp  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningFlow  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/core/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningApp\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningFlow\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/core/app.py:15\u001b[0m\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _console\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrequest_types\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRequest, CommandRequest, DeltaRequest\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconstants\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m     17\u001b[0m     DEBUG_ENABLED,\n\u001b[1;32m     18\u001b[0m     FLOW_DURATION_SAMPLES,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m     STATE_ACCUMULATE_WAIT,\n\u001b[1;32m     22\u001b[0m )\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/api/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mhttp_methods\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Delete, Get, Post, Put\n\u001b[1;32m      3\u001b[0m __all__ \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDelete\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGet\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPost\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPut\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/api/http_methods.py:10\u001b[0m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01muuid\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m uuid4\n\u001b[0;32m---> 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FastAPI, HTTPException\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrequest_types\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRequest, CommandRequest, RequestResponse\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/__init__.py:5\u001b[0m\n\u001b[1;32m      3\u001b[0m __version__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m0.1.17\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FastAPI\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrouting\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRouter\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/applications.py:3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Dict, List, Optional, Type\n\u001b[0;32m----> 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m routing\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mopenapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdocs\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_redoc_html, get_swagger_ui_html\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/routing.py:6\u001b[0m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, List, Optional, Type\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m params\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdependencies\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodels\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dependant\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/params.py:4\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Sequence\n\u001b[0;32m----> 4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpydantic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Schema\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mParamTypes\u001b[39;00m(Enum):\n",
-      "\u001b[0;31mImportError\u001b[0m: cannot import name 'Schema' from 'pydantic' (/home/aktersnurra/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/pydantic/__init__.py)",
-      "\nDuring handling of the above exception, another exception occurred:\n",
-      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[1], line 13\u001b[0m\n\u001b[1;32m     10\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01msys\u001b[39;00m\n\u001b[1;32m     11\u001b[0m     sys\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mappend(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124m..\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtext_recognizer\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01memnist\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m EMNIST\n",
-      "File \u001b[0;32m~/projects/text-recognizer/text_recognizer/data/emnist.py:14\u001b[0m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtoml\u001b[39;00m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mloguru\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m logger \u001b[38;5;28;01mas\u001b[39;00m log\n\u001b[0;32m---> 14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtext_recognizer\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase_data_module\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseDataModule, load_and_print_info\n\u001b[1;32m     15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtext_recognizer\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase_dataset\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseDataset, split_dataset\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtext_recognizer\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdownload_utils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m download_dataset\n",
-      "File \u001b[0;32m~/projects/text-recognizer/text_recognizer/data/base_data_module.py:5\u001b[0m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpathlib\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Path\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Callable, Dict, Optional, Tuple, TypeVar\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mL\u001b[39;00m\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m DataLoader\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtext_recognizer\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdata\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbase_dataset\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseDataset\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning/__init__.py:31\u001b[0m\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m__about__\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;241m*\u001b[39m  \u001b[38;5;66;03m# noqa: E402, F401, F403\u001b[39;00m\n\u001b[1;32m     30\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m__version__\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m version \u001b[38;5;28;01mas\u001b[39;00m __version__  \u001b[38;5;66;03m# noqa: E402, F401\u001b[39;00m\n\u001b[0;32m---> 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m storage  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningApp  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningFlow  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning/app/__init__.py:37\u001b[0m\n\u001b[1;32m     34\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[1;32m     36\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mos\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m linesep\n\u001b[0;32m---> 37\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m __version__\n\u001b[1;32m     38\u001b[0m     msg \u001b[38;5;241m=\u001b[39m \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mYour `lightning` package was built for `lightning_app==1.8.0rc0`, but you are running \u001b[39m\u001b[38;5;132;01m{\u001b[39;00m__version__\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m'\u001b[39m\n\u001b[1;32m     39\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mtype\u001b[39m(err)(\u001b[38;5;28mstr\u001b[39m(err) \u001b[38;5;241m+\u001b[39m linesep \u001b[38;5;241m+\u001b[39m msg)\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/__init__.py:31\u001b[0m\n\u001b[1;32m     28\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(__about__, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m__version__\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[1;32m     29\u001b[0m     \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01m__version__\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m version \u001b[38;5;28;01mas\u001b[39;00m __version__  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n\u001b[0;32m---> 31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningApp  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningFlow  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n\u001b[1;32m     33\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mwork\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningWork  \u001b[38;5;66;03m# noqa: E402\u001b[39;00m\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/core/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningApp\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mflow\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningFlow\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mwork\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LightningWork\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/core/app.py:15\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_utilities\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapply_func\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m apply_to_collection\n\u001b[1;32m     14\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m _console\n\u001b[0;32m---> 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrequest_types\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRequest, CommandRequest, DeltaRequest\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconstants\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m     17\u001b[0m     DEBUG_ENABLED,\n\u001b[1;32m     18\u001b[0m     FLOW_DURATION_SAMPLES,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     21\u001b[0m     STATE_ACCUMULATE_WAIT,\n\u001b[1;32m     22\u001b[0m )\n\u001b[1;32m     23\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcore\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mqueues\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m BaseQueue, SingleProcessQueue\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/api/__init__.py:1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mhttp_methods\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Delete, Get, Post, Put\n\u001b[1;32m      3\u001b[0m __all__ \u001b[38;5;241m=\u001b[39m [\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mDelete\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mGet\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPost\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPut\u001b[39m\u001b[38;5;124m\"\u001b[39m]\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/lightning_app/api/http_methods.py:10\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Dict, List, Optional\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01muuid\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m uuid4\n\u001b[0;32m---> 10\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FastAPI, HTTPException\n\u001b[1;32m     12\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrequest_types\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRequest, CommandRequest, RequestResponse\n\u001b[1;32m     13\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mlightning_app\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutilities\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapp_helpers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Logger\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/__init__.py:5\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;124;03m\"\"\"FastAPI framework, high performance, easy to learn, fast to code, ready for production\"\"\"\u001b[39;00m\n\u001b[1;32m      3\u001b[0m __version__ \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m0.1.17\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mapplications\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FastAPI\n\u001b[1;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mrouting\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m APIRouter\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mparams\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Body, Path, Query, Header, Cookie, Form, File, Security, Depends\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/applications.py:3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Dict, List, Optional, Type\n\u001b[0;32m----> 3\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m routing\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mopenapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdocs\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_redoc_html, get_swagger_ui_html\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mopenapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_openapi\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/routing.py:6\u001b[0m\n\u001b[1;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mlogging\u001b[39;00m\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, List, Optional, Type\n\u001b[0;32m----> 6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m params\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdependencies\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mmodels\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dependant\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mfastapi\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mdependencies\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m get_body_field, get_dependant, solve_dependencies\n",
-      "File \u001b[0;32m~/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/fastapi/params.py:4\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01menum\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Enum\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtyping\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Any, Callable, Sequence\n\u001b[0;32m----> 4\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpydantic\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Schema\n\u001b[1;32m      7\u001b[0m \u001b[38;5;28;01mclass\u001b[39;00m \u001b[38;5;21;01mParamTypes\u001b[39;00m(Enum):\n\u001b[1;32m      8\u001b[0m     query \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquery\u001b[39m\u001b[38;5;124m\"\u001b[39m\n",
-      "\u001b[0;31mImportError\u001b[0m: cannot import name 'Schema' from 'pydantic' (/home/aktersnurra/.cache/pypoetry/virtualenvs/text-recognizer-ejNaVa9M-py3.9/lib/python3.9/site-packages/pydantic/__init__.py)"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "%load_ext autoreload\n",
     "%autoreload 2\n",
diff --git a/notebooks/03-look-at-iam-lines.ipynb b/notebooks/03-look-at-iam-lines.ipynb
index 3d71c3c..9e9b24c 100644
--- a/notebooks/03-look-at-iam-lines.ipynb
+++ b/notebooks/03-look-at-iam-lines.ipynb
@@ -64,12 +64,12 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/tmp/ipykernel_20508/3921102733.py:2: UserWarning: \n",
+      "/tmp/ipykernel_3955/3921102733.py:2: UserWarning: \n",
       "The version_base parameter is not specified.\n",
       "Please specify a compatability version level, or None.\n",
       "Will assume defaults for version 1.1\n",
       "  with initialize(config_path=str(path.parent)):\n",
-      "/tmp/ipykernel_20508/3921102733.py:2: UserWarning: \n",
+      "/tmp/ipykernel_3955/3921102733.py:2: UserWarning: \n",
       "The version_base parameter is not specified.\n",
       "Please specify a compatability version level, or None.\n",
       "Will assume defaults for version 1.1\n",
@@ -108,8 +108,8 @@
       "Input dims: (1, 56, 1024)\n",
       "Output dims: (89, 1)\n",
       "Train/val/test sizes: 10255, 1140, 1958\n",
-      "Train Batch x stats: (torch.Size([8, 1, 56, 1024]), torch.float32, tensor(0.), tensor(0.0388), tensor(0.1288), tensor(1.))\n",
-      "Train Batch y stats: (torch.Size([8, 89]), torch.int64, tensor(1), tensor(52))\n",
+      "Train Batch x stats: (torch.Size([8, 1, 56, 1024]), torch.float32, tensor(0.), tensor(0.0361), tensor(0.1067), tensor(1.))\n",
+      "Train Batch y stats: (torch.Size([8, 89]), torch.int64, tensor(1), tensor(56))\n",
       "Test Batch x stats: (torch.Size([8, 1, 56, 1024]), torch.float32, tensor(0.), tensor(0.0333), tensor(0.0951), tensor(0.8627))\n",
       "Test Batch y stats: (torch.Size([8, 89]), torch.int64, tensor(1), tensor(52))\n",
       "\n"
@@ -125,7 +125,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -134,7 +134,40 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 22,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(tensor([[[0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "          ...,\n",
+       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
+       "          [0., 0., 0.,  ..., 0., 0., 0.]]]),\n",
+       " tensor([ 1, 32, 27, 14, 33, 16, 21, 22, 27, 20, 40, 14, 33, 40, 33, 21, 18, 40,\n",
+       "         29, 31, 28, 29, 40, 21, 28, 25, 17, 22, 27, 20, 40, 33, 21, 18, 40, 15,\n",
+       "         14, 31, 20, 18, 45, 32,  2,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,\n",
+       "          3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,\n",
+       "          3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3]))"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "datamodule.data_val[16]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
    "metadata": {},
    "outputs": [
     {
@@ -143,7 +176,7 @@
        "'union would be prepared to reach'"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -157,6 +190,112 @@
   },
   {
    "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['<b>',\n",
+       " '<s>',\n",
+       " '<e>',\n",
+       " '<p>',\n",
+       " '0',\n",
+       " '1',\n",
+       " '2',\n",
+       " '3',\n",
+       " '4',\n",
+       " '5',\n",
+       " '6',\n",
+       " '7',\n",
+       " '8',\n",
+       " '9',\n",
+       " 'a',\n",
+       " 'b',\n",
+       " 'c',\n",
+       " 'd',\n",
+       " 'e',\n",
+       " 'f',\n",
+       " 'g',\n",
+       " 'h',\n",
+       " 'i',\n",
+       " 'j',\n",
+       " 'k',\n",
+       " 'l',\n",
+       " 'm',\n",
+       " 'n',\n",
+       " 'o',\n",
+       " 'p',\n",
+       " 'q',\n",
+       " 'r',\n",
+       " 's',\n",
+       " 't',\n",
+       " 'u',\n",
+       " 'v',\n",
+       " 'w',\n",
+       " 'x',\n",
+       " 'y',\n",
+       " 'z',\n",
+       " ' ',\n",
+       " '!',\n",
+       " '\"',\n",
+       " '#',\n",
+       " '&',\n",
+       " \"'\",\n",
+       " '(',\n",
+       " ')',\n",
+       " '*',\n",
+       " '+',\n",
+       " ',',\n",
+       " '-',\n",
+       " '.',\n",
+       " '/',\n",
+       " ':',\n",
+       " ';',\n",
+       " '?',\n",
+       " '\\n']"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "datamodule.tokenizer.mapping"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([ 1, 34, 27, 22, 28, 27, 40, 36, 28, 34, 25, 17, 40, 15, 18, 40, 29, 31,\n",
+       "        18, 29, 14, 31, 18, 17, 40, 33, 28, 40, 31, 18, 14, 16, 21,  2,  3,  3,\n",
+       "         3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,\n",
+       "         3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,\n",
+       "         3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3,  3])"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "x =dataset[0][1]\n",
+    "x["
+   ]
+  },
+  {
+   "cell_type": "code",
    "execution_count": 10,
    "metadata": {},
    "outputs": [
diff --git a/notebooks/Untitled.ipynb b/notebooks/Untitled.ipynb
index 7ea06ae..5dbda4c 100644
--- a/notebooks/Untitled.ipynb
+++ b/notebooks/Untitled.ipynb
@@ -38,6 +38,30 @@
   {
    "cell_type": "code",
    "execution_count": 3,
+   "id": "4660ff12-74d7-4c8a-a65b-2c0d90326e26",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "ename": "AttributeError",
+     "evalue": "module 'torch' has no attribute 'cudnn'",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
+      "Cell \u001b[0;32mIn[3], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mtorch\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcudnn\u001b[49m\n",
+      "\u001b[0;31mAttributeError\u001b[0m: module 'torch' has no attribute 'cudnn'"
+     ]
+    }
+   ],
+   "source": [
+    "torch.cudnn"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
    "id": "652cfb26",
    "metadata": {
     "tags": []
diff --git a/notebooks/Untitled1.ipynb b/notebooks/Untitled1.ipynb
index 06129a3..a2d6168 100644
--- a/notebooks/Untitled1.ipynb
+++ b/notebooks/Untitled1.ipynb
@@ -536,8 +536,132 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "id": "c0e733d8-c17d-46f5-b484-9c74e46d7308",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "from typing import Union\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "\n",
+    "def first_appearance(x: torch.Tensor, element: Union[int, float], dim: int = 1) -> torch.Tensor:\n",
+    "    \"\"\"Return indices of first appearance of element in x, collapsing along dim.\n",
+    "\n",
+    "    Based on https://discuss.pytorch.org/t/first-nonzero-index/24769/9\n",
+    "\n",
+    "    Parameters\n",
+    "    ----------\n",
+    "    x\n",
+    "        One or two-dimensional Tensor to search for element.\n",
+    "    element\n",
+    "        Item to search for inside x.\n",
+    "    dim\n",
+    "        Dimension of Tensor to collapse over.\n",
+    "\n",
+    "    Returns\n",
+    "    -------\n",
+    "    torch.Tensor\n",
+    "        Indices where element occurs in x. If element is not found,\n",
+    "        return length of x along dim. One dimension smaller than x.\n",
+    "\n",
+    "    Raises\n",
+    "    ------\n",
+    "    ValueError\n",
+    "        if x is not a 1 or 2 dimensional Tensor\n",
+    "\n",
+    "    Examples\n",
+    "    --------\n",
+    "    >>> first_appearance(torch.tensor([[1, 2, 3], [2, 3, 3], [1, 1, 1], [3, 1, 1]]), 3)\n",
+    "    tensor([2, 1, 3, 0])\n",
+    "    >>> first_appearance(torch.tensor([1, 2, 3]), 1, dim=0)\n",
+    "    tensor(0)\n",
+    "    \"\"\"\n",
+    "    if x.dim() > 2 or x.dim() == 0:\n",
+    "        raise ValueError(f\"only 1 or 2 dimensional Tensors allowed, got Tensor with dim {x.dim()}\")\n",
+    "    matches = x == element\n",
+    "    first_appearance_mask = (matches.cumsum(dim) == 1) & matches\n",
+    "    does_match, match_index = first_appearance_mask.max(dim)\n",
+    "    first_inds = torch.where(does_match, match_index, x.shape[dim])\n",
+    "    return first_inds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "26ff2314-b2df-408f-b83a-f5fc903145da",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([2, 3])"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "first_appearance(torch.tensor([[1, 1, 3], [1, 1, 1]]), 3)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "e8c9dd16-4917-40bc-8504-084035882ced",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([ True, False])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "torch.any(torch.isin(torch.tensor([[1, 1, 3], [1, 1, 1]]), 3), 1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "30eb1926-53d2-431a-b7c1-f95919887b84",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "tensor([2, 0])"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "torch.tensor([[1, 1, 3], [1, 1, 1]]).argmax(dim=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a02c61b3-c84f-4778-90d0-fe6aafa12ccc",
    "metadata": {},
    "outputs": [],
    "source": []