diff --git a/04-custom-train-retrieval.ipynb b/04-custom-train-retrieval.ipynb
index ccd5357..676c426 100644
--- a/04-custom-train-retrieval.ipynb
+++ b/04-custom-train-retrieval.ipynb
@@ -18,7 +18,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 4,
    "id": "03855d9e-afe2-4edc-8b78-6b1d44b109ce",
    "metadata": {},
    "outputs": [
@@ -40,7 +40,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 5,
    "id": "7322bf4e-0bb0-486c-b041-6882ecee7239",
    "metadata": {},
    "outputs": [
@@ -118,7 +118,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 6,
    "id": "319b0b73-d1bc-4a90-a4aa-864d2215d646",
    "metadata": {},
    "outputs": [],
@@ -146,7 +146,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 7,
    "id": "2e4b6da8-8246-410c-8718-6a0dea4ea269",
    "metadata": {},
    "outputs": [],
@@ -166,7 +166,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 8,
    "id": "3325fcee-3bd1-4c1c-8ff8-5fac4be8fc16",
    "metadata": {},
    "outputs": [
@@ -176,7 +176,7 @@
        "'vocab_dict.pkl'"
       ]
      },
-     "execution_count": 5,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -187,7 +187,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 9,
    "id": "03bbe4be-68ed-4084-b4cd-eee9a311da58",
    "metadata": {},
    "outputs": [],
@@ -199,7 +199,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 10,
    "id": "d9dfaf03-9b67-438d-bcd6-131a75d9d51d",
    "metadata": {},
    "outputs": [
@@ -236,7 +236,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 11,
    "id": "03410e35-7e1c-481c-aef2-55c25a90bcc4",
    "metadata": {},
    "outputs": [
@@ -299,7 +299,7 @@
      "output_type": "stream",
      "text": [
       "EXPERIMENT_NAME: scale-training-v1\n",
-      "RUN_NAME: run-20230922-202508\n"
+      "RUN_NAME: run-20230925-145451\n"
      ]
     }
    ],
@@ -330,7 +330,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "TB_RESOURCE_NAME: projects/934903580331/locations/us-central1/tensorboards/4282210762094542848\n",
+      "TB_RESOURCE_NAME: projects/934903580331/locations/us-central1/tensorboards/7204483970304573440\n",
       "TB display name: scale-training-v1\n"
      ]
     }
@@ -565,7 +565,7 @@
       "                              '--candidate_file_dir=ndr-v1-hybrid-vertex-bucket',\n",
       "                              '--candidate_files_prefix=data/v1/candidates',\n",
       "                              '--experiment_name=scale-training-v1',\n",
-      "                              '--experiment_run=run-20230922-202508',\n",
+      "                              '--experiment_run=run-20230925-145451',\n",
       "                              '--num_epochs=15',\n",
       "                              '--batch_size=4096',\n",
       "                              '--embedding_dim=128',\n",
@@ -580,7 +580,7 @@
       "                              '--pipeline_version=v1',\n",
       "                              '--seed=1234',\n",
       "                              '--max_tokens=20000',\n",
-      "                              '--tb_resource_name=projects/934903580331/locations/us-central1/tensorboards/4282210762094542848',\n",
+      "                              '--tb_resource_name=projects/934903580331/locations/us-central1/tensorboards/7204483970304573440',\n",
       "                              '--embed_frequency=1',\n",
       "                              '--hist_frequency=0',\n",
       "                              '--tf_gpu_thread_count=8',\n",
@@ -639,7 +639,7 @@
      "output_type": "stream",
      "text": [
       "\n",
-      " Copied training package and Dockerfile to gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230922-202508\n",
+      " Copied training package and Dockerfile to gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-145451\n",
       "\n"
      ]
     }
@@ -668,9 +668,9 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230922-202508/Dockerfile_tfrs\n",
-      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230922-202508/vocab_dict.pkl\n",
-      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230922-202508/trainer/\n"
+      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-145451/Dockerfile_tfrs\n",
+      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-145451/vocab_dict.pkl\n",
+      "gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-145451/trainer/\n"
      ]
     }
    ],
@@ -696,7 +696,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "JOB_NAME: train-v1-run-20230922-202508\n"
+      "JOB_NAME: train-v1-run-20230925-145451\n"
      ]
     }
    ],
@@ -766,7 +766,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 12,
    "id": "33a58184-e0c4-4403-ad80-da5fdb7dbaed",
    "metadata": {},
    "outputs": [
@@ -774,7 +774,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "TB_LOGS_PATH: gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230922-153317/logs\n"
+      "TB_LOGS_PATH: gs://ndr-v1-hybrid-vertex-bucket/scale-training-v1/run-20230925-145451/logs\n"
      ]
     }
    ],
@@ -789,7 +789,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": 14,
    "id": "d432d552-4416-4f4c-964f-bcca57b60a07",
    "metadata": {},
    "outputs": [],
@@ -800,29 +800,20 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 15,
    "id": "171227db-d5f9-47e4-8b67-a528950233ab",
    "metadata": {},
    "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Reusing TensorBoard on port 6006 (pid 30025), started 0:04:41 ago. (Use '!kill 30025' to kill it.)"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
     {
      "data": {
       "text/html": [
        "\n",
-       "      <iframe id=\"tensorboard-frame-64233e23ab6abec0\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
+       "      <iframe id=\"tensorboard-frame-633ff50d405cd4f5\" width=\"100%\" height=\"800\" frameborder=\"0\">\n",
        "      </iframe>\n",
        "      <script>\n",
        "        (function() {\n",
-       "          const frame = document.getElementById(\"tensorboard-frame-64233e23ab6abec0\");\n",
-       "          const url = new URL(\"/proxy/6006/\", window.location);\n",
+       "          const frame = document.getElementById(\"tensorboard-frame-633ff50d405cd4f5\");\n",
+       "          const url = new URL(\"/proxy/6007/\", window.location);\n",
        "          const port = 0;\n",
        "          if (port) {\n",
        "            url.port = port;\n",
diff --git a/06-matching-engine.ipynb b/06-deploy-index-and-query-tower.ipynb
similarity index 100%
rename from 06-matching-engine.ipynb
rename to 06-deploy-index-and-query-tower.ipynb
diff --git a/07-train-pipeline.ipynb b/07-train-pipeline.ipynb
index 0ea8071..204883a 100644
--- a/07-train-pipeline.ipynb
+++ b/07-train-pipeline.ipynb
@@ -234,7 +234,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "PIPELINE_VERSION = 'pipev1'       # TODO"
+    "PIPELINE_VERSION = 'pipev2'       # TODO"
    ]
   },
   {
@@ -1416,7 +1416,7 @@
     "\n",
     "    index_endpoint = index_endpoint.deploy_index(\n",
     "        index=brute_index, \n",
-    "        deployed_index_id=f'{deployed_brute_force_index_name}', #-{TIMESTAMP}'\n",
+    "        deployed_index_id=f'{deployed_brute_force_index_name}' #-{TIMESTAMP}'\n",
     "    )\n",
     "\n",
     "    logging.info(f\"index_endpoint.deployed_indexes: {index_endpoint.deployed_indexes}\")\n",
@@ -1466,7 +1466,7 @@
     "    base_image=\"python:3.9\",\n",
     "    packages_to_install=[\n",
     "        'google-cloud-aiplatform==1.26.1',\n",
-    "        # 'google-cloud-pipeline-components',\n",
+    "        'google-cloud-pipeline-components',\n",
     "        'google-cloud-storage',\n",
     "        'tensorflow==2.11.0',\n",
     "        'numpy'\n",
@@ -1478,6 +1478,7 @@
     "    version: str,\n",
     "    prefix: str,\n",
     "    emails: str,\n",
+    "    train_output_gcs_bucket: str,\n",
     "    # feature_dict: dict, # TODO\n",
     "    bq_dataset: str,\n",
     "    bq_train_table: str,\n",
@@ -1566,7 +1567,7 @@
     "    # ===================================================\n",
     "    QUERY_FILENAME = 'query_feats_dict.pkl'\n",
     "    # FEATURES_PREFIX = f'{experiment_name}/{experiment_run}/features'\n",
-    "    GCS_PATH_TO_BLOB = f'{experiment_name}/{experiment_run}/features'\n",
+    "    GCS_PATH_TO_BLOB = f'{experiment_name}/{experiment_run}/features/{QUERY_FILENAME}'\n",
     "    \n",
     "    loaded_feat_dict = download_blob(\n",
     "        bucket_name=train_output_gcs_bucket,\n",
@@ -1575,7 +1576,7 @@
     "    )\n",
     "    logging.info(f'loaded_feat_dict: {loaded_feat_dict}')\n",
     "    \n",
-    "    filehandler = open(loaded_feat_dict, 'rb')\n",
+    "    filehandler = open(QUERY_FILENAME, 'rb')\n",
     "    FEAT_DICT = pkl.load(filehandler)\n",
     "    filehandler.close()\n",
     "    \n",
@@ -1654,7 +1655,7 @@
     "        display_name=JOB_DISPLAY_NAME,\n",
     "        project=project,\n",
     "        location=location,\n",
-    "        endpoint=endpoint,\n",
+    "        endpoint=_endpoint,\n",
     "        logging_sampling_strategy=logging_sampling_strategy,\n",
     "        schedule_config=schedule_config,\n",
     "        alert_config=alert_config,\n",
@@ -1699,7 +1700,7 @@
     "    base_image=\"python:3.9\",\n",
     "    packages_to_install=[\n",
     "        'google-cloud-aiplatform==1.26.1',\n",
-    "        # 'google-cloud-pipeline-components',\n",
+    "        'google-cloud-pipeline-components',\n",
     "        'google-cloud-storage',\n",
     "        'tensorflow==2.11.0',\n",
     "        'numpy'\n",
@@ -2218,7 +2219,8 @@
    "source": [
     "EXPERIMENT_PREFIX = 'tfrs-pipe'                     # custom identifier for organizing experiments\n",
     "EXPERIMENT_NAME=f'{EXPERIMENT_PREFIX}-{VERSION}'\n",
-    "RUN_NAME = f'run-{time.strftime(\"%Y%m%d-%H%M%S\")}'\n",
+    "# RUN_NAME = f'run-{time.strftime(\"%Y%m%d-%H%M%S\")}'\n",
+    "RUN_NAME = f'run-20230922-202528'\n",
     "\n",
     "print(f\"EXPERIMENT_NAME: {EXPERIMENT_NAME}\")\n",
     "print(f\"RUN_NAME: {RUN_NAME}\")"
@@ -2558,7 +2560,7 @@
     "valid_samples_cnt    = 82_959\n",
     "\n",
     "# validation & evaluation\n",
-    "VALID_FREQUENCY      = NUM_EPOCHS // 2 # 20\n",
+    "VALID_FREQUENCY      = NUM_EPOCHS // 3 # 20\n",
     "VALID_STEPS          = valid_samples_cnt // BATCH_SIZE # 100\n",
     "EPOCH_STEPS          = train_sample_cnt // BATCH_SIZE\n",
     "\n",
@@ -3123,6 +3125,7 @@
     "            version=model_version,\n",
     "            prefix=prefix,\n",
     "            emails=emails,\n",
+    "            train_output_gcs_bucket=train_output_gcs_bucket,\n",
     "            # feature_dict=feature_dict, # TODO\n",
     "            bq_dataset= bq_dataset,\n",
     "            bq_train_table=bq_train_table,\n",
@@ -3207,7 +3210,7 @@
     "            version=model_version,\n",
     "            location=location,\n",
     "            vpc_network_name=vpc_network_name,\n",
-    "            ann_index_endpoint_display_name=f'ann-index-endpoint_{model_version}'.replace('-', '_'),\n",
+    "            ann_index_endpoint_display_name=f'ann-index-endpoint_{pipeline_version}'.replace('-', '_'),\n",
     "            ann_index_endpoint_description='endpoint for ann index',\n",
     "            ann_index_resource_uri=create_ann_index_op.outputs['ann_index_resource_uri'],\n",
     "        )\n",
@@ -3223,7 +3226,7 @@
     "            version=model_version,\n",
     "            location=location,\n",
     "            vpc_network_name=vpc_network_name,\n",
-    "            brute_index_endpoint_display_name=f'bf-index-endpoint_{model_version}'.replace('-', '_'),\n",
+    "            brute_index_endpoint_display_name=f'bf-index-endpoint_{pipeline_version}'.replace('-', '_'),\n",
     "            brute_index_endpoint_description='endpoint for brute force index',\n",
     "            brute_force_index_resource_uri=create_brute_force_index_op.outputs['brute_force_index_resource_uri'],\n",
     "        )\n",
@@ -3348,6 +3351,7 @@
      "text": [
       "gs://ndr-v1-hybrid-vertex-bucket/tfrs-pipe-v1/run-20230922-202528/pipeline_root/Dockerfile_tfrs\n",
       "gs://ndr-v1-hybrid-vertex-bucket/tfrs-pipe-v1/run-20230922-202528/pipeline_root/pipeline_spec.json\n",
+      "gs://ndr-v1-hybrid-vertex-bucket/tfrs-pipe-v1/run-20230922-202528/pipeline_root/934903580331/\n",
       "gs://ndr-v1-hybrid-vertex-bucket/tfrs-pipe-v1/run-20230922-202528/pipeline_root/trainer/\n"
      ]
     }
diff --git a/08-build-ranking-model.ipynb b/08-build-ranking-model.ipynb
index c056ee8..96cf060 100644
--- a/08-build-ranking-model.ipynb
+++ b/08-build-ranking-model.ipynb
@@ -10,7 +10,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 26,
    "id": "e742fb79-7e4f-4ef2-b60d-938b1828b667",
    "metadata": {},
    "outputs": [
@@ -37,7 +37,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 27,
    "id": "034e35b3-824f-4b4e-8b47-84fb8ed1410f",
    "metadata": {},
    "outputs": [
@@ -59,7 +59,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 28,
    "id": "0ea84ba3-a0f5-45ad-a075-3118d19c4833",
    "metadata": {},
    "outputs": [
@@ -145,7 +145,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 29,
    "id": "a0a913e4-cd52-4dab-838b-e85df80f7209",
    "metadata": {},
    "outputs": [],
@@ -160,7 +160,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 30,
    "id": "67bf2da7-a8df-4ee7-9b36-cc423a8cf0f5",
    "metadata": {},
    "outputs": [],
@@ -187,20 +187,20 @@
     "# ranking models\n",
     "from util import feature_set_utils as feature_utils\n",
     "from src.ranking import tf_ranking_model as tfrm\n",
+    "\n",
     "# from src.ranking import build_audio_ranker_v1 as build_audio_ranker\n",
     "from src.ranking import train_utils\n",
+    "\n",
     "# feature_sets\n",
     "\n",
     "# GPU\n",
     "import gc\n",
-    "from numba import cuda\n",
-    "\n",
-    "# jw-repo2/spotify_mpd_two_tower/src/ranking/train_utils.py"
+    "from numba import cuda"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 33,
    "id": "1f264ea3-7267-4439-b4bb-140b683c3826",
    "metadata": {},
    "outputs": [
@@ -222,7 +222,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 35,
    "id": "dd6a0af0-cb71-400c-8a8d-e0c1b5fd21e2",
    "metadata": {},
    "outputs": [
@@ -232,7 +232,7 @@
        "37"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 35,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -253,7 +253,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 36,
    "id": "ef62eb48-d062-4eff-859e-ba52b4b455b2",
    "metadata": {},
    "outputs": [
@@ -263,7 +263,13 @@
      "text": [
       "\u001b[01;34msrc/ranking\u001b[00m\n",
       "├── __init__.py\n",
-      "├── build_audio_ranker_v1.py\n",
+      "├── \u001b[01;34m__pycache__\u001b[00m\n",
+      "│   ├── __init__.cpython-37.pyc\n",
+      "│   ├── build_audio_ranker.cpython-37.pyc\n",
+      "│   ├── tf_ranking_model.cpython-37.pyc\n",
+      "│   ├── train_config.cpython-37.pyc\n",
+      "│   └── train_utils.cpython-37.pyc\n",
+      "├── build_audio_ranker.py\n",
       "├── feature_sets.py\n",
       "├── requirements.txt\n",
       "├── task.py\n",
@@ -271,7 +277,7 @@
       "├── train_config.py\n",
       "└── train_utils.py\n",
       "\n",
-      "0 directories, 8 files\n"
+      "1 directory, 13 files\n"
      ]
     }
    ],
@@ -289,7 +295,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 38,
    "id": "a6ae7a7c-454d-487c-862f-983efe3ccd73",
    "metadata": {},
    "outputs": [],
@@ -310,7 +316,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 45,
+   "execution_count": 39,
    "id": "f0695cad-c232-4638-a304-bd7aae3ebdb2",
    "metadata": {},
    "outputs": [],
@@ -352,7 +358,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 40,
    "id": "e8ac8961-e942-4f93-89d6-2949274c1029",
    "metadata": {
     "tags": []
@@ -374,10 +380,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 41,
    "id": "c7dcc62a-7ab8-4b37-8fa0-98e3236bc6d1",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "ename": "TypeError",
+     "evalue": "in user code:\n\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/train_utils.py\", line 107, in full_parse  *\n        data = tf.data.TFRecordDataset(data)\n\n    TypeError: The `filenames` argument must contain `tf.string` elements. Got `tf.float32` elements.\n",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[0;31mTypeError\u001b[0m                                 Traceback (most recent call last)",
+      "\u001b[0;32m~/jw-repo2/spotify_mpd_two_tower/ipykernel_24177/3079836612.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     33\u001b[0m     \u001b[0mnum_parallel_calls\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAUTOTUNE\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     34\u001b[0m     \u001b[0mcycle_length\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtf\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mAUTOTUNE\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 35\u001b[0;31m     \u001b[0mdeterministic\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     36\u001b[0m \u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     37\u001b[0m     \u001b[0mfeature_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparse_rank_tfrecord\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;31m#parse_rank_tfrecord,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py\u001b[0m in \u001b[0;36minterleave\u001b[0;34m(self, map_func, cycle_length, block_length, num_parallel_calls, deterministic, name)\u001b[0m\n\u001b[1;32m   2491\u001b[0m           \u001b[0mnum_parallel_calls\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2492\u001b[0m           \u001b[0mdeterministic\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mdeterministic\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2493\u001b[0;31m           name=name)\n\u001b[0m\u001b[1;32m   2494\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2495\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0mfilter\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mpredicate\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/dataset_ops.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, input_dataset, map_func, cycle_length, block_length, num_parallel_calls, buffer_output_elements, prefetch_input_elements, deterministic, name)\u001b[0m\n\u001b[1;32m   5667\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_input_dataset\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0minput_dataset\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5668\u001b[0m     self._map_func = structured_function.StructuredFunctionWrapper(\n\u001b[0;32m-> 5669\u001b[0;31m         map_func, self._transformation_name(), dataset=input_dataset)\n\u001b[0m\u001b[1;32m   5670\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_map_func\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0moutput_structure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mDatasetSpec\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   5671\u001b[0m       raise TypeError(\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/structured_function.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, func, transformation_name, dataset, input_classes, input_shapes, input_types, input_structure, add_to_graph, use_legacy_function, defun_kwargs)\u001b[0m\n\u001b[1;32m    261\u001b[0m         \u001b[0mfn_factory\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrace_tf_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdefun_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    262\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 263\u001b[0;31m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn_factory\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    264\u001b[0m     \u001b[0;31m# There is no graph to add in eager mode.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    265\u001b[0m     \u001b[0madd_to_graph\u001b[0m \u001b[0;34m&=\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mcontext\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexecuting_eagerly\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py\u001b[0m in \u001b[0;36mget_concrete_function\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    225\u001b[0m     \"\"\"\n\u001b[1;32m    226\u001b[0m     concrete_function = self._get_concrete_function_garbage_collected(\n\u001b[0;32m--> 227\u001b[0;31m         *args, **kwargs)\n\u001b[0m\u001b[1;32m    228\u001b[0m     \u001b[0mconcrete_function\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_garbage_collector\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrelease\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    229\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mconcrete_function\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py\u001b[0m in \u001b[0;36m_get_concrete_function_garbage_collected\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m    190\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    191\u001b[0m     \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_lock\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 192\u001b[0;31m       \u001b[0mconcrete_function\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_define_concrete_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    193\u001b[0m       \u001b[0mseen_names\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    194\u001b[0m       captured = object_identity.ObjectIdentitySet(\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py\u001b[0m in \u001b[0;36m_maybe_define_concrete_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m    155\u001b[0m       \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m{\u001b[0m\u001b[0;34m}\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    156\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 157\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_maybe_define_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    158\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    159\u001b[0m   \u001b[0;32mdef\u001b[0m \u001b[0m_get_concrete_function_internal_garbage_collected\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py\u001b[0m in \u001b[0;36m_maybe_define_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m    358\u001b[0m             \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgeneralized_func_key\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_placeholder_value\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    359\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 360\u001b[0;31m           \u001b[0mconcrete_function\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_create_concrete_function\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    361\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    362\u001b[0m           \u001b[0mgraph_capture_container\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconcrete_function\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgraph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_capture_func_lib\u001b[0m  \u001b[0;31m# pylint: disable=protected-access\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/polymorphic_function/tracing_compiler.py\u001b[0m in \u001b[0;36m_create_concrete_function\u001b[0;34m(self, args, kwargs)\u001b[0m\n\u001b[1;32m    291\u001b[0m             \u001b[0mautograph_options\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_autograph_options\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    292\u001b[0m             \u001b[0marg_names\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0marg_names\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 293\u001b[0;31m             capture_by_value=self._capture_by_value),\n\u001b[0m\u001b[1;32m    294\u001b[0m         \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_function_attributes\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    295\u001b[0m         \u001b[0mspec\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfunction_spec\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/framework/func_graph.py\u001b[0m in \u001b[0;36mfunc_graph_from_py_func\u001b[0;34m(name, python_func, args, kwargs, signature, func_graph, autograph, autograph_options, add_control_dependencies, arg_names, op_return_value, collections, capture_by_value, acd_record_initial_resource_uses)\u001b[0m\n\u001b[1;32m   1281\u001b[0m         \u001b[0m_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moriginal_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtf_decorator\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0munwrap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpython_func\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1282\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1283\u001b[0;31m       \u001b[0mfunc_outputs\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpython_func\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mfunc_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mfunc_kwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1284\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1285\u001b[0m       \u001b[0;31m# invariant: `func_outputs` contains only Tensors, CompositeTensors,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/structured_function.py\u001b[0m in \u001b[0;36mwrapped_fn\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    238\u001b[0m           attributes=defun_kwargs)\n\u001b[1;32m    239\u001b[0m       \u001b[0;32mdef\u001b[0m \u001b[0mwrapped_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint: disable=missing-docstring\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 240\u001b[0;31m         \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mwrapper_helper\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    241\u001b[0m         \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mstructure\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_tensor_list\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_output_structure\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mret\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    242\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0mops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert_to_tensor\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mt\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mt\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mret\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/structured_function.py\u001b[0m in \u001b[0;36mwrapper_helper\u001b[0;34m(*args)\u001b[0m\n\u001b[1;32m    169\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0m_should_unpack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnested_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    170\u001b[0m         \u001b[0mnested_args\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mnested_args\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 171\u001b[0;31m       \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mautograph\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtf_convert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_func\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mag_ctx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0mnested_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    172\u001b[0m       \u001b[0mret\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mvariable_utils\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert_variables_to_tensors\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mret\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    173\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0m_should_pack\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mret\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    690\u001b[0m       \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint:disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    691\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'ag_error_metadata'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 692\u001b[0;31m           \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mag_error_metadata\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto_exception\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    693\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    694\u001b[0m           \u001b[0;32mraise\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py\u001b[0m in \u001b[0;36mwrapper\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    687\u001b[0m       \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    688\u001b[0m         \u001b[0;32mwith\u001b[0m \u001b[0mconversion_ctx\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 689\u001b[0;31m           \u001b[0;32mreturn\u001b[0m \u001b[0mconverted_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0moptions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    690\u001b[0m       \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m  \u001b[0;31m# pylint:disable=broad-except\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    691\u001b[0m         \u001b[0;32mif\u001b[0m \u001b[0mhasattr\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'ag_error_metadata'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py\u001b[0m in \u001b[0;36mconverted_call\u001b[0;34m(f, args, kwargs, caller_fn_scope, options)\u001b[0m\n\u001b[1;32m    437\u001b[0m     \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    438\u001b[0m       \u001b[0;32mif\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 439\u001b[0;31m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconverted_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0meffective_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    440\u001b[0m       \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    441\u001b[0m         \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconverted_f\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0meffective_args\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/jw-repo2/spotify_mpd_two_tower/src/ranking/train_utils.py\u001b[0m in \u001b[0;36mtf__full_parse\u001b[0;34m(data)\u001b[0m\n\u001b[1;32m      9\u001b[0m \u001b[0;31m# ====================================================\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     10\u001b[0m \u001b[0;31m# Helper functions\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 11\u001b[0;31m \u001b[0;31m# ====================================================\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     12\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     13\u001b[0m \u001b[0;31m# set project # env var\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py\u001b[0m in \u001b[0;36mconverted_call\u001b[0;34m(f, args, kwargs, caller_fn_scope, options)\u001b[0m\n\u001b[1;32m    329\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mconversion\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mis_in_allowlist_cache\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    330\u001b[0m     \u001b[0mlogging\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mlog\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m2\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'Allowlisted %s: from cache'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 331\u001b[0;31m     \u001b[0;32mreturn\u001b[0m \u001b[0m_call_unconverted\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptions\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    332\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    333\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mag_ctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontrol_status_ctx\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstatus\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0mag_ctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mStatus\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mDISABLED\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/autograph/impl/api.py\u001b[0m in \u001b[0;36m_call_unconverted\u001b[0;34m(f, args, kwargs, options, update_cache)\u001b[0m\n\u001b[1;32m    457\u001b[0m   \u001b[0;32mif\u001b[0m \u001b[0mkwargs\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    458\u001b[0m     \u001b[0;32mreturn\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 459\u001b[0;31m   \u001b[0;32mreturn\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    460\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    461\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/readers.py\u001b[0m in \u001b[0;36m__init__\u001b[0;34m(self, filenames, compression_type, buffer_size, num_parallel_reads, name)\u001b[0m\n\u001b[1;32m    457\u001b[0m       \u001b[0mValueError\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mIf\u001b[0m \u001b[0many\u001b[0m \u001b[0margument\u001b[0m \u001b[0mdoes\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mhave\u001b[0m \u001b[0mthe\u001b[0m \u001b[0mexpected\u001b[0m \u001b[0mshape\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    458\u001b[0m     \"\"\"\n\u001b[0;32m--> 459\u001b[0;31m     \u001b[0mfilenames\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_create_or_validate_filenames_dataset\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilenames\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mname\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    460\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    461\u001b[0m     \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_filenames\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfilenames\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/data/ops/readers.py\u001b[0m in \u001b[0;36m_create_or_validate_filenames_dataset\u001b[0;34m(filenames, name)\u001b[0m\n\u001b[1;32m     68\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mfilenames\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mdtypes\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstring\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m       raise TypeError(\n\u001b[0;32m---> 70\u001b[0;31m           \u001b[0;34m\"The `filenames` argument must contain `tf.string` elements. Got \"\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     71\u001b[0m           f\"`{filenames.dtype!r}` elements.\")\n\u001b[1;32m     72\u001b[0m     \u001b[0mfilenames\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0marray_ops\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mreshape\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfilenames\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"flat_filenames\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;31mTypeError\u001b[0m: in user code:\n\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/train_utils.py\", line 107, in full_parse  *\n        data = tf.data.TFRecordDataset(data)\n\n    TypeError: The `filenames` argument must contain `tf.string` elements. Got `tf.float32` elements.\n"
+     ]
+    }
+   ],
    "source": [
     "valid_files = []\n",
     "\n",
@@ -434,7 +471,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "id": "ffa3c07b-dc77-4c54-8b8b-74d410970ddd",
    "metadata": {},
    "outputs": [
@@ -462,7 +499,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 14,
    "id": "08db584f-6f81-466a-abc1-5b49f5a18894",
    "metadata": {},
    "outputs": [],
@@ -474,7 +511,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 15,
    "id": "8d675e1c-700a-4dc6-b91f-c5c1a15be08e",
    "metadata": {},
    "outputs": [
@@ -511,7 +548,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 16,
    "id": "f1c62297-0fde-49ae-9ab9-ab0666c16d9a",
    "metadata": {},
    "outputs": [
@@ -519,8 +556,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "RUN_NAME: run-20230922-210502\n",
-      "LOG_DIR: gs://ndr-v1-hybrid-vertex-bucket/local-ranker-v1/run-20230922-210502/tb-logs/\n"
+      "RUN_NAME: run-20230925-163646\n",
+      "LOG_DIR: gs://ndr-v1-hybrid-vertex-bucket/local-ranker-v1/run-20230925-163646/tb-logs/\n"
      ]
     }
    ],
@@ -558,7 +595,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 17,
    "id": "70ad6de9-5650-4561-821f-3c889209a9b9",
    "metadata": {},
    "outputs": [
@@ -567,11 +604,11 @@
      "output_type": "stream",
      "text": [
       "Creating Tensorboard\n",
-      "Create Tensorboard backing LRO: projects/934903580331/locations/us-central1/tensorboards/5225151934075240448/operations/1313984221315858432\n",
-      "Tensorboard created. Resource name: projects/934903580331/locations/us-central1/tensorboards/5225151934075240448\n",
+      "Create Tensorboard backing LRO: projects/934903580331/locations/us-central1/tensorboards/5899003028320550912/operations/5118567132304506880\n",
+      "Tensorboard created. Resource name: projects/934903580331/locations/us-central1/tensorboards/5899003028320550912\n",
       "To use this Tensorboard in another session:\n",
-      "tb = aiplatform.Tensorboard('projects/934903580331/locations/us-central1/tensorboards/5225151934075240448')\n",
-      "TB_RESOURCE_NAME: projects/934903580331/locations/us-central1/tensorboards/5225151934075240448\n"
+      "tb = aiplatform.Tensorboard('projects/934903580331/locations/us-central1/tensorboards/5899003028320550912')\n",
+      "TB_RESOURCE_NAME: projects/934903580331/locations/us-central1/tensorboards/5899003028320550912\n"
      ]
     }
    ],
@@ -598,17 +635,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 18,
    "id": "6108f696-9a14-4aae-ac69-9ed9603f5ad2",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from src.ranking import build_audio_ranker_v3 as build_audio_ranker"
+    "from src.ranking import build_audio_ranker as build_audio_ranker"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": 19,
    "id": "49c6d5ea-261e-4837-a1ec-fb451735694d",
    "metadata": {},
    "outputs": [],
@@ -632,7 +669,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": 20,
    "id": "f4f7818a-03f1-4678-b476-a51a10978b55",
    "metadata": {},
    "outputs": [],
@@ -668,7 +705,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": 21,
    "id": "b9f388d5-894b-4887-855d-c8353a24201c",
    "metadata": {},
    "outputs": [],
@@ -678,23 +715,32 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 31,
-   "id": "188edf7b-64b5-4cc0-86cc-198ed4c540f6",
+   "execution_count": 22,
+   "id": "3233d69e-0e3a-434a-9cc8-0a9a95f88f4f",
    "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "<src.ranking.build_audio_ranker_v3.TheRankingModel at 0x7f76b40c2550>"
-      ]
-     },
-     "execution_count": 31,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
+   "outputs": [],
+   "source": [
+    "# model.summary()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "id": "9d96ce0b-c1f8-4047-b9d8-04152793254b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model.ranking_model.summary()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "bf5afd75-2faf-4f57-b48d-432a13484c69",
+   "metadata": {},
+   "outputs": [],
    "source": [
-    "model"
+    "# model.fit(valid_dataset, epochs=1) # valid_dataset | train_dataset"
    ]
   },
   {
@@ -707,7 +753,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 32,
+   "execution_count": 33,
    "id": "060115a4-0dd9-4cf1-b4fc-85c3e9123fc1",
    "metadata": {},
    "outputs": [
@@ -766,7 +812,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 33,
+   "execution_count": 34,
    "id": "646035ea-deb8-4f64-a457-fe7958c72b2a",
    "metadata": {},
    "outputs": [],
@@ -774,7 +820,7 @@
     "NUM_EPOCHS = 3\n",
     "VALID_FREQUENCY = 5\n",
     "HIST_FREQ = 0\n",
-    "EMBED_FREQ = 1\n",
+    "EMBED_FREQ = 0\n",
     "\n",
     "LOCAL_TRAIN_DIR = f\"local_train_dir/{EXPERIMENT_NAME}/{RUN_NAME}/\"\n",
     "LOCAL_CHECKPOINT_DIR = f\"{LOCAL_TRAIN_DIR}/chkpts\" # my_model.ckpt\n",
@@ -796,15 +842,15 @@
     },
     {
      "ename": "InvalidArgumentError",
-     "evalue": "Graph execution error:\n\nDetected at node 'ranking/mean_squared_error/SquaredDifference' defined at (most recent call last):\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 193, in _run_module_as_main\n      \"__main__\", mod_spec)\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 85, in _run_code\n      exec(code, run_globals)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py\", line 17, in <module>\n      app.launch_new_instance()\n    File \"/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py\", line 1041, in launch_instance\n      app.start()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py\", line 712, in start\n      self.io_loop.start()\n    File \"/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py\", line 215, in start\n      self.asyncio_loop.run_forever()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 541, in run_forever\n      self._run_once()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 1786, in _run_once\n      handle._run()\n    File \"/opt/conda/lib/python3.7/asyncio/events.py\", line 88, in _run\n      self._context.run(self._callback, *self._args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 510, in dispatch_queue\n      await self.process_one()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 499, in process_one\n      await dispatch(*args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 406, in dispatch_shell\n      await result\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 730, in execute_request\n      reply_content = await reply_content\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/ipkernel.py\", line 387, in do_execute\n      cell_id=cell_id,\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py\", line 528, in run_cell\n      return super().run_cell(*args, **kwargs)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 2976, in run_cell\n      raw_cell, store_history, silent, shell_futures, cell_id\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n      return runner(coro)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n      coro.send(None)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3258, in run_cell_async\n      interactivity=interactivity, compiler=compiler, result=result)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n      if (await self.run_code(code, result,  async_=asy)):\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n      exec(code_obj, self.user_global_ns, self.user_ns)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/ipykernel_319/299002834.py\", line 41, in <module>\n      verbose=1\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1650, in fit\n      tmp_logs = self.train_function(iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1249, in train_function\n      return step_function(self, iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1233, in step_function\n      outputs = model.distribute_strategy.run(run_step, args=(data,))\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1222, in run_step\n      outputs = model.train_step(data)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/models/base.py\", line 68, in train_step\n      loss = self.compute_loss(inputs, training=True)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/build_audio_ranker_v3.py\", line 602, in compute_loss\n      , predictions=tf.squeeze(rating_predictions, axis=-1),\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/base_layer.py\", line 1132, in __call__\n      outputs = call_fn(inputs, *args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 96, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/tasks/ranking.py\", line 89, in call\n      loss = self._loss(\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 152, in __call__\n      losses = call_fn(y_true, y_pred)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 284, in call\n      return ag_fn(y_true, y_pred, **self._fn_kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 1500, in mean_squared_error\n      return backend.mean(tf.math.squared_difference(y_pred, y_true), axis=-1)\nNode: 'ranking/mean_squared_error/SquaredDifference'\nrequired broadcastable shapes\n\t [[{{node ranking/mean_squared_error/SquaredDifference}}]] [Op:__inference_train_function_7488]",
+     "evalue": "Graph execution error:\n\nDetected at node 'ranking/mean_squared_error/SquaredDifference' defined at (most recent call last):\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 193, in _run_module_as_main\n      \"__main__\", mod_spec)\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 85, in _run_code\n      exec(code, run_globals)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py\", line 17, in <module>\n      app.launch_new_instance()\n    File \"/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py\", line 1041, in launch_instance\n      app.start()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py\", line 712, in start\n      self.io_loop.start()\n    File \"/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py\", line 215, in start\n      self.asyncio_loop.run_forever()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 541, in run_forever\n      self._run_once()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 1786, in _run_once\n      handle._run()\n    File \"/opt/conda/lib/python3.7/asyncio/events.py\", line 88, in _run\n      self._context.run(self._callback, *self._args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 510, in dispatch_queue\n      await self.process_one()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 499, in process_one\n      await dispatch(*args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 406, in dispatch_shell\n      await result\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 730, in execute_request\n      reply_content = await reply_content\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/ipkernel.py\", line 387, in do_execute\n      cell_id=cell_id,\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py\", line 528, in run_cell\n      return super().run_cell(*args, **kwargs)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 2976, in run_cell\n      raw_cell, store_history, silent, shell_futures, cell_id\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n      return runner(coro)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n      coro.send(None)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3258, in run_cell_async\n      interactivity=interactivity, compiler=compiler, result=result)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n      if (await self.run_code(code, result,  async_=asy)):\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n      exec(code_obj, self.user_global_ns, self.user_ns)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/ipykernel_22821/299002834.py\", line 41, in <module>\n      verbose=1\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1650, in fit\n      tmp_logs = self.train_function(iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1249, in train_function\n      return step_function(self, iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1233, in step_function\n      outputs = model.distribute_strategy.run(run_step, args=(data,))\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1222, in run_step\n      outputs = model.train_step(data)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/models/base.py\", line 68, in train_step\n      loss = self.compute_loss(inputs, training=True)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/build_audio_ranker.py\", line 615, in compute_loss\n      , predictions=tf.squeeze(rating_predictions, axis=-1),\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/base_layer.py\", line 1132, in __call__\n      outputs = call_fn(inputs, *args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 96, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/tasks/ranking.py\", line 89, in call\n      loss = self._loss(\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 152, in __call__\n      losses = call_fn(y_true, y_pred)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 284, in call\n      return ag_fn(y_true, y_pred, **self._fn_kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 1500, in mean_squared_error\n      return backend.mean(tf.math.squared_difference(y_pred, y_true), axis=-1)\nNode: 'ranking/mean_squared_error/SquaredDifference'\nrequired broadcastable shapes\n\t [[{{node ranking/mean_squared_error/SquaredDifference}}]] [Op:__inference_train_function_6407]",
      "output_type": "error",
      "traceback": [
       "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
       "\u001b[0;31mInvalidArgumentError\u001b[0m                      Traceback (most recent call last)",
-      "\u001b[0;32m~/jw-repo2/spotify_mpd_two_tower/ipykernel_319/299002834.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     39\u001b[0m         \u001b[0mmodel_checkpoint_callback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     ], \n\u001b[0;32m---> 41\u001b[0;31m     \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     42\u001b[0m )\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
+      "\u001b[0;32m~/jw-repo2/spotify_mpd_two_tower/ipykernel_22821/299002834.py\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     39\u001b[0m         \u001b[0mmodel_checkpoint_callback\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     40\u001b[0m     ], \n\u001b[0;32m---> 41\u001b[0;31m     \u001b[0mverbose\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     42\u001b[0m )\n\u001b[1;32m     43\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;32m~/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\u001b[0m in \u001b[0;36merror_handler\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m     68\u001b[0m             \u001b[0;31m# To get the full stack trace, call:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     69\u001b[0m             \u001b[0;31m# `tf.debugging.disable_traceback_filtering()`\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 70\u001b[0;31m             \u001b[0;32mraise\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mwith_traceback\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfiltered_tb\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     71\u001b[0m         \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     72\u001b[0m             \u001b[0;32mdel\u001b[0m \u001b[0mfiltered_tb\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;32m~/.local/lib/python3.7/site-packages/tensorflow/python/eager/execute.py\u001b[0m in \u001b[0;36mquick_execute\u001b[0;34m(op_name, num_outputs, inputs, attrs, ctx, name)\u001b[0m\n\u001b[1;32m     51\u001b[0m     \u001b[0mctx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mensure_initialized\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m     tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,\n\u001b[0;32m---> 53\u001b[0;31m                                         inputs, attrs, num_outputs)\n\u001b[0m\u001b[1;32m     54\u001b[0m   \u001b[0;32mexcept\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_NotOkStatusException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0mname\u001b[0m \u001b[0;32mis\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
-      "\u001b[0;31mInvalidArgumentError\u001b[0m: Graph execution error:\n\nDetected at node 'ranking/mean_squared_error/SquaredDifference' defined at (most recent call last):\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 193, in _run_module_as_main\n      \"__main__\", mod_spec)\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 85, in _run_code\n      exec(code, run_globals)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py\", line 17, in <module>\n      app.launch_new_instance()\n    File \"/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py\", line 1041, in launch_instance\n      app.start()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py\", line 712, in start\n      self.io_loop.start()\n    File \"/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py\", line 215, in start\n      self.asyncio_loop.run_forever()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 541, in run_forever\n      self._run_once()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 1786, in _run_once\n      handle._run()\n    File \"/opt/conda/lib/python3.7/asyncio/events.py\", line 88, in _run\n      self._context.run(self._callback, *self._args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 510, in dispatch_queue\n      await self.process_one()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 499, in process_one\n      await dispatch(*args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 406, in dispatch_shell\n      await result\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 730, in execute_request\n      reply_content = await reply_content\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/ipkernel.py\", line 387, in do_execute\n      cell_id=cell_id,\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py\", line 528, in run_cell\n      return super().run_cell(*args, **kwargs)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 2976, in run_cell\n      raw_cell, store_history, silent, shell_futures, cell_id\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n      return runner(coro)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n      coro.send(None)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3258, in run_cell_async\n      interactivity=interactivity, compiler=compiler, result=result)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n      if (await self.run_code(code, result,  async_=asy)):\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n      exec(code_obj, self.user_global_ns, self.user_ns)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/ipykernel_319/299002834.py\", line 41, in <module>\n      verbose=1\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1650, in fit\n      tmp_logs = self.train_function(iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1249, in train_function\n      return step_function(self, iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1233, in step_function\n      outputs = model.distribute_strategy.run(run_step, args=(data,))\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1222, in run_step\n      outputs = model.train_step(data)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/models/base.py\", line 68, in train_step\n      loss = self.compute_loss(inputs, training=True)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/build_audio_ranker_v3.py\", line 602, in compute_loss\n      , predictions=tf.squeeze(rating_predictions, axis=-1),\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/base_layer.py\", line 1132, in __call__\n      outputs = call_fn(inputs, *args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 96, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/tasks/ranking.py\", line 89, in call\n      loss = self._loss(\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 152, in __call__\n      losses = call_fn(y_true, y_pred)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 284, in call\n      return ag_fn(y_true, y_pred, **self._fn_kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 1500, in mean_squared_error\n      return backend.mean(tf.math.squared_difference(y_pred, y_true), axis=-1)\nNode: 'ranking/mean_squared_error/SquaredDifference'\nrequired broadcastable shapes\n\t [[{{node ranking/mean_squared_error/SquaredDifference}}]] [Op:__inference_train_function_7488]"
+      "\u001b[0;31mInvalidArgumentError\u001b[0m: Graph execution error:\n\nDetected at node 'ranking/mean_squared_error/SquaredDifference' defined at (most recent call last):\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 193, in _run_module_as_main\n      \"__main__\", mod_spec)\n    File \"/opt/conda/lib/python3.7/runpy.py\", line 85, in _run_code\n      exec(code, run_globals)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel_launcher.py\", line 17, in <module>\n      app.launch_new_instance()\n    File \"/opt/conda/lib/python3.7/site-packages/traitlets/config/application.py\", line 1041, in launch_instance\n      app.start()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelapp.py\", line 712, in start\n      self.io_loop.start()\n    File \"/opt/conda/lib/python3.7/site-packages/tornado/platform/asyncio.py\", line 215, in start\n      self.asyncio_loop.run_forever()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 541, in run_forever\n      self._run_once()\n    File \"/opt/conda/lib/python3.7/asyncio/base_events.py\", line 1786, in _run_once\n      handle._run()\n    File \"/opt/conda/lib/python3.7/asyncio/events.py\", line 88, in _run\n      self._context.run(self._callback, *self._args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 510, in dispatch_queue\n      await self.process_one()\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 499, in process_one\n      await dispatch(*args)\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 406, in dispatch_shell\n      await result\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/kernelbase.py\", line 730, in execute_request\n      reply_content = await reply_content\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/ipkernel.py\", line 387, in do_execute\n      cell_id=cell_id,\n    File \"/opt/conda/lib/python3.7/site-packages/ipykernel/zmqshell.py\", line 528, in run_cell\n      return super().run_cell(*args, **kwargs)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 2976, in run_cell\n      raw_cell, store_history, silent, shell_futures, cell_id\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3030, in _run_cell\n      return runner(coro)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/async_helpers.py\", line 78, in _pseudo_sync_runner\n      coro.send(None)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3258, in run_cell_async\n      interactivity=interactivity, compiler=compiler, result=result)\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3473, in run_ast_nodes\n      if (await self.run_code(code, result,  async_=asy)):\n    File \"/opt/conda/lib/python3.7/site-packages/IPython/core/interactiveshell.py\", line 3553, in run_code\n      exec(code_obj, self.user_global_ns, self.user_ns)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/ipykernel_22821/299002834.py\", line 41, in <module>\n      verbose=1\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1650, in fit\n      tmp_logs = self.train_function(iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1249, in train_function\n      return step_function(self, iterator)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1233, in step_function\n      outputs = model.distribute_strategy.run(run_step, args=(data,))\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/training.py\", line 1222, in run_step\n      outputs = model.train_step(data)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/models/base.py\", line 68, in train_step\n      loss = self.compute_loss(inputs, training=True)\n    File \"/home/jupyter/jw-repo2/spotify_mpd_two_tower/src/ranking/build_audio_ranker.py\", line 615, in compute_loss\n      , predictions=tf.squeeze(rating_predictions, axis=-1),\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 65, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/engine/base_layer.py\", line 1132, in __call__\n      outputs = call_fn(inputs, *args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/utils/traceback_utils.py\", line 96, in error_handler\n      return fn(*args, **kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/tensorflow_recommenders/tasks/ranking.py\", line 89, in call\n      loss = self._loss(\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 152, in __call__\n      losses = call_fn(y_true, y_pred)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 284, in call\n      return ag_fn(y_true, y_pred, **self._fn_kwargs)\n    File \"/home/jupyter/.local/lib/python3.7/site-packages/keras/losses.py\", line 1500, in mean_squared_error\n      return backend.mean(tf.math.squared_difference(y_pred, y_true), axis=-1)\nNode: 'ranking/mean_squared_error/SquaredDifference'\nrequired broadcastable shapes\n\t [[{{node ranking/mean_squared_error/SquaredDifference}}]] [Op:__inference_train_function_6407]"
      ]
     }
    ],
diff --git a/src/ranking/build_audio_ranker.py b/src/ranking/build_audio_ranker.py
index e72fafb..1f1ceeb 100644
--- a/src/ranking/build_audio_ranker.py
+++ b/src/ranking/build_audio_ranker.py
@@ -565,13 +565,26 @@ def __init__(
             , embedding_dim = embedding_dim
             , projection_dim = projection_dim
         )
-
-        self.task = tfrs.tasks.Ranking(
-          loss=tf.keras.losses.MeanSquaredError(),
-          metrics=[
-            tf.keras.metrics.RootMeanSquaredError()
-          ]
+        
+        self.task: tf.keras.layers.Layer = tfrs.tasks.Ranking(
+            loss = tf.keras.losses.MeanSquaredError(
+                # reduction=tf.keras.losses.Reduction.SUM
+            )
+            , metrics = [
+                tf.keras.metrics.RootMeanSquaredError(name="rmse_metric")
+                # , tfr.keras.metrics.NDCGMetric(
+                #     name="ndcg_metric"
+                #     , ragged = True
+                # )
+            ]
         )
+
+        # self.task = tfrs.tasks.Ranking(
+        #   loss=tf.keras.losses.MeanSquaredError(),
+        #   metrics=[
+        #     tf.keras.metrics.RootMeanSquaredError()
+        #   ]
+        # )
             
         # tfr.keras.metrics.NDCGMetric(name="ndcg_metric"),
         #     loss = tf.keras.losses.MeanSquaredError(
diff --git a/src/train_pipes/deploy_brute_index.py b/src/train_pipes/deploy_brute_index.py
index 01ef446..e936aed 100644
--- a/src/train_pipes/deploy_brute_index.py
+++ b/src/train_pipes/deploy_brute_index.py
@@ -48,7 +48,7 @@ def deploy_brute_index(
 
     index_endpoint = index_endpoint.deploy_index(
         index=brute_index, 
-        deployed_index_id=f'{deployed_brute_force_index_name}', #-{TIMESTAMP}'
+        deployed_index_id=f'{deployed_brute_force_index_name}' #-{TIMESTAMP}'
     )
 
     logging.info(f"index_endpoint.deployed_indexes: {index_endpoint.deployed_indexes}")
diff --git a/src/train_pipes/model_monitoring_config.py b/src/train_pipes/model_monitoring_config.py
index 532e46e..c39127c 100644
--- a/src/train_pipes/model_monitoring_config.py
+++ b/src/train_pipes/model_monitoring_config.py
@@ -9,7 +9,7 @@
     base_image="python:3.9",
     packages_to_install=[
         'google-cloud-aiplatform==1.26.1',
-        # 'google-cloud-pipeline-components',
+        'google-cloud-pipeline-components',
         'google-cloud-storage',
         'tensorflow==2.11.0',
         'numpy'
@@ -21,6 +21,7 @@ def model_monitoring_config(
     version: str,
     prefix: str,
     emails: str,
+    train_output_gcs_bucket: str,
     # feature_dict: dict, # TODO
     bq_dataset: str,
     bq_train_table: str,
@@ -109,7 +110,7 @@ def download_blob(bucket_name, source_gcs_obj, local_filename):
     # ===================================================
     QUERY_FILENAME = 'query_feats_dict.pkl'
     # FEATURES_PREFIX = f'{experiment_name}/{experiment_run}/features'
-    GCS_PATH_TO_BLOB = f'{experiment_name}/{experiment_run}/features'
+    GCS_PATH_TO_BLOB = f'{experiment_name}/{experiment_run}/features/{QUERY_FILENAME}'
     
     loaded_feat_dict = download_blob(
         bucket_name=train_output_gcs_bucket,
@@ -118,7 +119,7 @@ def download_blob(bucket_name, source_gcs_obj, local_filename):
     )
     logging.info(f'loaded_feat_dict: {loaded_feat_dict}')
     
-    filehandler = open(loaded_feat_dict, 'rb')
+    filehandler = open(QUERY_FILENAME, 'rb')
     FEAT_DICT = pkl.load(filehandler)
     filehandler.close()
     
@@ -197,7 +198,7 @@ def download_blob(bucket_name, source_gcs_obj, local_filename):
         display_name=JOB_DISPLAY_NAME,
         project=project,
         location=location,
-        endpoint=endpoint,
+        endpoint=_endpoint,
         logging_sampling_strategy=logging_sampling_strategy,
         schedule_config=schedule_config,
         alert_config=alert_config,
diff --git a/src/train_pipes/test_model_endpoint.py b/src/train_pipes/test_model_endpoint.py
index 9805f5f..36f90db 100644
--- a/src/train_pipes/test_model_endpoint.py
+++ b/src/train_pipes/test_model_endpoint.py
@@ -9,7 +9,7 @@
     base_image="python:3.9",
     packages_to_install=[
         'google-cloud-aiplatform==1.26.1',
-        # 'google-cloud-pipeline-components',
+        'google-cloud-pipeline-components',
         'google-cloud-storage',
         'tensorflow==2.11.0',
         'numpy'