fix: metrics correction

notebooks/03_fine_tuning.ipynb

@@ -26,7 +26,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
       "metadata": {},
       "outputs": [
         {
@@ -66,7 +66,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 2,
       "metadata": {},
       "outputs": [
         {
@@ -112,7 +112,7 @@
         "    # General\n",
         "    \"seed\": 42,\n",
         "    \"device\": \"cuda\" if torch.cuda.is_available() else \"cpu\",\n",
-        "    \"num_workers\": 4,  # Number of worker processes for DataLoader\n",
+        "    \"num_workers\": 0,  # Number of worker processes for DataLoader\n",
         "}\n",
         "\n",
         "# Set random seed\n",
@@ -132,27 +132,9 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 2,
+      "execution_count": null,
       "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "--2025-12-09 18:33:50--  https://ftp.ncbi.nlm.nih.gov/genomes/refseq/vertebrate_mammalian/Homo_sapiens/latest_assembly_versions/GCF_000001405.40_GRCh38.p14/GCF_000001405.40_GRCh38.p14_genomic.fna.gz\n",
-            "Resolving ftp.ncbi.nlm.nih.gov (ftp.ncbi.nlm.nih.gov)... 2607:f220:41e:250::7, 2607:f220:41e:250::11, 2607:f220:41e:250::12, ...\n",
-            "Connecting to ftp.ncbi.nlm.nih.gov (ftp.ncbi.nlm.nih.gov)|2607:f220:41e:250::7|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 972898531 (928M) [application/x-gzip]\n",
-            "Saving to: 'GCF_000001405.40_GRCh38.p14_genomic.fna.gz'\n",
-            "\n",
-            "GCF_000001405.40_GR 100%[===================>] 927.83M  18.4MB/s    in 51s     \n",
-            "\n",
-            "2025-12-09 18:34:42 (18.0 MB/s) - 'GCF_000001405.40_GRCh38.p14_genomic.fna.gz' saved [972898531/972898531]\n",
-            "\n"
-          ]
-        }
-      ],
+      "outputs": [],
       "source": [
         "!wget -c https://ftp.ncbi.nlm.nih.gov/genomes/refseq/vertebrate_mammalian/Homo_sapiens/latest_assembly_versions/GCF_000001405.40_GRCh38.p14/GCF_000001405.40_GRCh38.p14_genomic.fna.gz \\\n",
         "&& gunzip -f GCF_000001405.40_GRCh38.p14_genomic.fna.gz"
@@ -160,22 +142,9 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 16,
+      "execution_count": null,
       "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "--2025-12-09 22:13:59--  https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL\n",
-            "Resolving www.encodeproject.org (www.encodeproject.org)... 34.211.244.144\n",
-            "Connecting to www.encodeproject.org (www.encodeproject.org)|34.211.244.144|:443... connected.\n",
-            "HTTP request sent, awaiting response... 404 Not Found\n",
-            "2025-12-09 22:14:00 ERROR 404: Not Found.\n",
-            "\n"
-          ]
-        }
-      ],
+      "outputs": [],
       "source": [
         "!wget -O ENCFF884LDL \"$(curl -s https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL | sed -n 's/.*href=\\\"\\([^\\\"]*ENCFF884LDL[^\\\"]*\\)\\\".*/\\1/p')\" \\\n",
         "&& echo \"Downloaded ENCFF884LDL\""
@@ -183,39 +152,16 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 4,
+      "execution_count": null,
       "metadata": {},
-      "outputs": [
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "--2025-12-09 18:41:24--  https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL.bigWig\n",
-            "Resolving www.encodeproject.org (www.encodeproject.org)... 34.211.244.144\n",
-            "Connecting to www.encodeproject.org (www.encodeproject.org)|34.211.244.144|:443... connected.\n",
-            "HTTP request sent, awaiting response... 307 Temporary Redirect\n",
-            "Location: https://encode-public.s3.amazonaws.com/2020/09/19/425880b6-b323-4ee2-95ce-56bdd088d126/ENCFF884LDL.bigWig?response-content-disposition=attachment%3B%20filename%3DENCFF884LDL.bigWig&AWSAccessKeyId=ASIATGZNGCNX3AXUNFS3&Signature=Ca%2Bz1PL7zdbGzyRggtvN686q4oE%3D&x-amz-security-token=IQoJb3JpZ2luX2VjEPr%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCXVzLXdlc3QtMiJGMEQCIAggXesBwHBuGSivVx0RvF5f2vZbk09TPBdf%2FYJUt%2BLWAiAKrh58c%2Bm%2F%2ByrujtQxgltFGzGo5qXSWv%2B0zPaa3gKUTCq8BQjC%2F%2F%2F%2F%2F%2F%2F%2F%2F%2F8BEAAaDDIyMDc0ODcxNDg2MyIMa%2FegIMq%2By2ql10quKpAFATT6r6oWCSXqrBd2gfR8S1QNvY%2BKjvbr%2BvS2ifnF5NqfByJgZxdXVC65WI8fUYqgspTQB5Az%2BE5O4jR8EnFBv%2FjO6DqrWkQQOUsHUFFGXJjarvCPdYjqJmV9SyeTuzNeV0xwFX%2Fleq1%2F4f3eAV81Nv5J%2B8UeHYn5GxtwS%2BjhzVsCJ8tqAo6yRi0wPteU8nb8yLJb%2F%2FWvQLZce7Yc9%2BZkuxKKGoEKQstRSGLCh%2FjtnNfvGp0x20mj5C7wsk61LHBJlNV3KVD7qZHZ57N1CBx5XNuJ%2BkJp6eBU8htM%2FY73tBkp4w5xHNyI5F%2B7JxjDDjo4YOikyLKk7tnTmWfC2lEGXXx33D8xyBxi4oNnK76R0N296GRSHS22esmo12YGK5QNvVbU4SuZUUWjVcrGFqtN%2F7ff1K%2FdqiRyh6TDvXbOUf%2Bk691iqwRY34LbXoJsOzcux5wwQGbHfcSdGrp2Y3KtpDGEdHiiTVHJeHi9pxBvlwvmjM5lXjJjtjOFqXIF%2F%2FygXdl4wUIMMsuinPWpA5xVIk4kg1Bv5XVNuqcPJl7Dl2ZdRzQvwc0Xl5dBL39ZAz9MvCffPV2Fb3hiL5vIQJ2ySdDnqXDhTuUsWGy81MltoznoOVbvuu64FAEp4GdwnwRH1ILlVOKQ1bHR5FSHqb8OFVqAQezRljaJY2ds1J2HMAJ2AJtg3k8XNQScR%2FutxWkI3pYDnAQQQkHHw3aFWNNYbQMfyAAptJohtNGClRoTiepBUckqxpgvMXwEOTJzpUEi0sMIxMkXMWa3ncKFHQAP6P3eKxBOjW8s%2F3BXwRlbgsNdQvqDUdf2dD5KLeHfpyKbdPnG0C6yZAxBF%2Fk4jO1F2F4o533RZGF8Ww7qMc5Ij2ww%2BbPhyQY6sgG2uZfWDKxd1yRNOufiZW%2FAtmcEQg%2BtzoWnq6TxyhU0OCY%2BN7xR8HO4UaT0Od0C06PHugNQCUS6eJusR0IfSRJ7ozZJUomphTeCPXw1G%2B6RVsni%2B9lGE8SlRLTMzNvzQJv8oJNZsoi6DVWlK%2FGt7TgwxSKH8%2BVQmal7nXUqR9f8Dh7CF1KppbVtNiGDaxTIN%2F7j%2BwIFrKHIMOYhC1dt5gPFnIQwnj1%2BuyEw5FWF3hKIkD%2Bc&Expires=1765431685 [following]\n",
-            "--2025-12-09 18:41:25--  https://encode-public.s3.amazonaws.com/2020/09/19/425880b6-b323-4ee2-95ce-56bdd088d126/ENCFF884LDL.bigWig?response-content-disposition=attachment%3B%20filename%3DENCFF884LDL.bigWig&AWSAccessKeyId=ASIATGZNGCNX3AXUNFS3&Signature=Ca%2Bz1PL7zdbGzyRggtvN686q4oE%3D&x-amz-security-token=IQoJb3JpZ2luX2VjEPr%2F%2F%2F%2F%2F%2F%2F%2F%2F%2FwEaCXVzLXdlc3QtMiJGMEQCIAggXesBwHBuGSivVx0RvF5f2vZbk09TPBdf%2FYJUt%2BLWAiAKrh58c%2Bm%2F%2ByrujtQxgltFGzGo5qXSWv%2B0zPaa3gKUTCq8BQjC%2F%2F%2F%2F%2F%2F%2F%2F%2F%2F8BEAAaDDIyMDc0ODcxNDg2MyIMa%2FegIMq%2By2ql10quKpAFATT6r6oWCSXqrBd2gfR8S1QNvY%2BKjvbr%2BvS2ifnF5NqfByJgZxdXVC65WI8fUYqgspTQB5Az%2BE5O4jR8EnFBv%2FjO6DqrWkQQOUsHUFFGXJjarvCPdYjqJmV9SyeTuzNeV0xwFX%2Fleq1%2F4f3eAV81Nv5J%2B8UeHYn5GxtwS%2BjhzVsCJ8tqAo6yRi0wPteU8nb8yLJb%2F%2FWvQLZce7Yc9%2BZkuxKKGoEKQstRSGLCh%2FjtnNfvGp0x20mj5C7wsk61LHBJlNV3KVD7qZHZ57N1CBx5XNuJ%2BkJp6eBU8htM%2FY73tBkp4w5xHNyI5F%2B7JxjDDjo4YOikyLKk7tnTmWfC2lEGXXx33D8xyBxi4oNnK76R0N296GRSHS22esmo12YGK5QNvVbU4SuZUUWjVcrGFqtN%2F7ff1K%2FdqiRyh6TDvXbOUf%2Bk691iqwRY34LbXoJsOzcux5wwQGbHfcSdGrp2Y3KtpDGEdHiiTVHJeHi9pxBvlwvmjM5lXjJjtjOFqXIF%2F%2FygXdl4wUIMMsuinPWpA5xVIk4kg1Bv5XVNuqcPJl7Dl2ZdRzQvwc0Xl5dBL39ZAz9MvCffPV2Fb3hiL5vIQJ2ySdDnqXDhTuUsWGy81MltoznoOVbvuu64FAEp4GdwnwRH1ILlVOKQ1bHR5FSHqb8OFVqAQezRljaJY2ds1J2HMAJ2AJtg3k8XNQScR%2FutxWkI3pYDnAQQQkHHw3aFWNNYbQMfyAAptJohtNGClRoTiepBUckqxpgvMXwEOTJzpUEi0sMIxMkXMWa3ncKFHQAP6P3eKxBOjW8s%2F3BXwRlbgsNdQvqDUdf2dD5KLeHfpyKbdPnG0C6yZAxBF%2Fk4jO1F2F4o533RZGF8Ww7qMc5Ij2ww%2BbPhyQY6sgG2uZfWDKxd1yRNOufiZW%2FAtmcEQg%2BtzoWnq6TxyhU0OCY%2BN7xR8HO4UaT0Od0C06PHugNQCUS6eJusR0IfSRJ7ozZJUomphTeCPXw1G%2B6RVsni%2B9lGE8SlRLTMzNvzQJv8oJNZsoi6DVWlK%2FGt7TgwxSKH8%2BVQmal7nXUqR9f8Dh7CF1KppbVtNiGDaxTIN%2F7j%2BwIFrKHIMOYhC1dt5gPFnIQwnj1%2BuyEw5FWF3hKIkD%2Bc&Expires=1765431685\n",
-            "Resolving encode-public.s3.amazonaws.com (encode-public.s3.amazonaws.com)... 3.5.81.13, 52.92.211.217, 52.92.197.57, ...\n",
-            "Connecting to encode-public.s3.amazonaws.com (encode-public.s3.amazonaws.com)|3.5.81.13|:443... connected.\n",
-            "HTTP request sent, awaiting response... 200 OK\n",
-            "Length: 568139478 (542M) [binary/octet-stream]\n",
-            "Saving to: 'ENCFF884LDL.bigWig'\n",
-            "\n",
-            "ENCFF884LDL.bigWig  100%[===================>] 541.82M  9.64MB/s    in 79s     \n",
-            "\n",
-            "2025-12-09 18:42:45 (6.88 MB/s) - 'ENCFF884LDL.bigWig' saved [568139478/568139478]\n",
-            "\n"
-          ]
-        }
-      ],
+      "outputs": [],
       "source": [
         "!wget -c https://www.encodeproject.org/files/ENCFF884LDL/@@download/ENCFF884LDL.bigWig"
       ]
     },
     {
       "cell_type": "code",
-      "execution_count": 5,
+      "execution_count": 3,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -265,7 +211,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 71,
+      "execution_count": 4,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -341,7 +287,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 72,
+      "execution_count": 5,
       "metadata": {},
       "outputs": [
         {
@@ -387,7 +333,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 34,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -500,20 +446,22 @@
         "        # For a single input string, its shape is typically (1, len(seq))\n",
         "\n",
         "        # Signal from bigWig tracks (numpy array) -> torch tensor\n",
-        "        bigwig_targets = [\n",
+        "        bigwig_targets = np.array([\n",
         "            self.bw_list[i].values(chrom, start, end, numpy=True)\n",
         "            for i in range(len(self.bw_list))\n",
-        "        ]\n",
+        "        ])  # shape (num_tracks, seq_len)\n",
+        "        # Transpose to (seq_len, num_tracks)\n",
+        "        bigwig_targets = bigwig_targets.T\n",
         "        # pyBigWig returns NaN where no data; turn NaN into 0\n",
         "        bigwig_targets = torch.tensor(bigwig_targets, dtype=torch.float32)\n",
         "        bigwig_targets = torch.nan_to_num(bigwig_targets, nan=0.0)\n",
         "        \n",
         "        # Crop targets to center fraction\n",
         "        if self.keep_target_center_fraction < 1.0:\n",
-        "            seq_len = bigwig_targets.shape[0]\n",
+        "            seq_len = bigwig_targets.shape[0]  # First dimension is sequence length\n",
         "            target_offset = int(seq_len * (1 - self.keep_target_center_fraction) // 2)\n",
         "            target_length = seq_len - 2 * target_offset\n",
-        "            bigwig_targets = bigwig_targets[target_offset:target_offset + target_length]\n",
+        "            bigwig_targets = bigwig_targets[target_offset:target_offset + target_length, :]\n",
         "\n",
         "        sample = {\n",
         "            \"tokens\": tokens,\n",
@@ -527,7 +475,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 35,
       "metadata": {},
       "outputs": [
         {
@@ -535,7 +483,8 @@
           "output_type": "stream",
           "text": [
             "Train samples: 100\n",
-            "Val samples: 10\n"
+            "Val samples: 10\n",
+            "Test samples: 10\n"
           ]
         }
       ],
@@ -605,7 +554,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 59,
+      "execution_count": 36,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -633,7 +582,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 60,
+      "execution_count": 37,
       "metadata": {},
       "outputs": [
         {
@@ -727,7 +676,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 38,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -794,20 +743,20 @@
         "        # Scaled metrics: per-track Pearson correlations\n",
         "        for i, (track_name, metric) in enumerate(zip(self.track_names, self.pearson_metrics_scaled)):\n",
         "            corr = metric.compute().item()\n",
-        "            metrics_dict[f\"{track_name}/pearson_scaled\"] = corr\n",
+        "            metrics_dict[f\"metrics_scaled/{track_name}/pearson\"] = corr\n",
         "        \n",
         "        # Scaled metrics: mean Pearson correlation\n",
         "        correlations_scaled = [metric.compute().item() for metric in self.pearson_metrics_scaled]\n",
-        "        metrics_dict[\"mean/pearson_scaled\"] = np.nanmean(correlations_scaled)\n",
+        "        metrics_dict[\"metrics_scaled/mean/pearson\"] = np.nanmean(correlations_scaled)\n",
         "        \n",
         "        # Raw metrics: per-track Pearson correlations\n",
         "        for i, (track_name, metric) in enumerate(zip(self.track_names, self.pearson_metrics_raw)):\n",
         "            corr = metric.compute().item()\n",
-        "            metrics_dict[f\"{track_name}/pearson_raw\"] = corr\n",
+        "            metrics_dict[f\"metrics_raw/{track_name}/pearson\"] = corr\n",
         "        \n",
         "        # Raw metrics: mean Pearson correlation\n",
         "        correlations_raw = [metric.compute().item() for metric in self.pearson_metrics_raw]\n",
-        "        metrics_dict[\"mean/pearson_raw\"] = np.nanmean(correlations_raw)\n",
+        "        metrics_dict[\"metrics_raw/mean/pearson\"] = np.nanmean(correlations_raw)\n",
         "        \n",
         "        # Mean loss\n",
         "        metrics_dict[\"loss\"] = np.mean(self.losses) if self.losses else 0.0\n",
@@ -817,7 +766,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 39,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -835,7 +784,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 63,
+      "execution_count": 40,
       "metadata": {},
       "outputs": [
         {
@@ -971,7 +920,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 64,
+      "execution_count": 41,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -1047,7 +996,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 42,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -1133,7 +1082,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 43,
       "metadata": {},
       "outputs": [
         {
@@ -1142,78 +1091,63 @@
           "text": [
             "Starting training...\n",
             "Training for 32 steps with 2 gradient accumulation steps\n",
-            "\n"
-          ]
-        },
-        {
-          "name": "stderr",
-          "output_type": "stream",
-          "text": [
-            "/home/y-bornachot/venvs/ntv3-env/lib/python3.12/site-packages/torch/amp/autocast_mode.py:287: UserWarning: In CPU autocast, but the target dtype is not supported. Disabling autocast.\n",
-            "CPU Autocast only supports dtype of torch.bfloat16, torch.float16 currently.\n",
-            "  warnings.warn(error_message)\n"
-          ]
-        },
-        {
-          "name": "stdout",
-          "output_type": "stream",
-          "text": [
-            "Step 0/32 | Loss: 1.5993 | Mean Pearson: -0.0848 | LR: 1.17e-09 | Tokens: 4,096\n",
+            "\n",
+            "Step 1/32 | Loss: 0.7569 | Mean Pearson: -0.1473 | LR: 1.17e-09 | Tokens: 4,096\n",
             "\n",
             "Running validation at step 0...\n",
-            "  Validation Loss: 0.6607\n",
-            "  Validation Mean Pearson: -0.0054\n",
-            "    ENCFF884LDL/pearson: -0.0054\n",
-            "Step 2/32 | Loss: 0.3453 | Mean Pearson: -0.2111 | LR: 2.50e-09 | Tokens: 12,288\n",
-            "Step 4/32 | Loss: 1.0248 | Mean Pearson: -0.0197 | LR: 2.41e-09 | Tokens: 20,480\n",
+            "  Validation Loss: 1.0152\n",
+            "  Validation Mean Pearson: -0.0414\n",
+            "    ENCFF884LDL/pearson: -0.0414\n",
+            "Step 3/32 | Loss: 0.3793 | Mean Pearson: -0.0229 | LR: 2.50e-09 | Tokens: 12,288\n",
+            "Step 5/32 | Loss: 0.4111 | Mean Pearson: -0.1739 | LR: 2.41e-09 | Tokens: 20,480\n",
             "\n",
             "Running validation at step 4...\n",
-            "  Validation Loss: 0.5158\n",
-            "  Validation Mean Pearson: 0.0160\n",
-            "    ENCFF884LDL/pearson: 0.0160\n",
-            "Step 6/32 | Loss: 0.3720 | Mean Pearson: 0.0140 | LR: 2.32e-09 | Tokens: 28,672\n",
-            "Step 8/32 | Loss: 0.4894 | Mean Pearson: -0.0300 | LR: 2.23e-09 | Tokens: 36,864\n",
+            "  Validation Loss: 0.4801\n",
+            "  Validation Mean Pearson: 0.0120\n",
+            "    ENCFF884LDL/pearson: 0.0120\n",
+            "Step 7/32 | Loss: 0.3404 | Mean Pearson: -0.0191 | LR: 2.32e-09 | Tokens: 28,672\n",
+            "Step 9/32 | Loss: 0.3950 | Mean Pearson: 0.0090 | LR: 2.23e-09 | Tokens: 36,864\n",
             "\n",
             "Running validation at step 8...\n",
-            "  Validation Loss: 0.5024\n",
-            "  Validation Mean Pearson: -0.0443\n",
-            "    ENCFF884LDL/pearson: -0.0443\n",
-            "Step 10/32 | Loss: 0.4039 | Mean Pearson: -0.0783 | LR: 2.13e-09 | Tokens: 45,056\n",
-            "Step 12/32 | Loss: 0.4974 | Mean Pearson: 0.0227 | LR: 2.02e-09 | Tokens: 53,248\n",
+            "  Validation Loss: 0.5865\n",
+            "  Validation Mean Pearson: -0.0260\n",
+            "    ENCFF884LDL/pearson: -0.0260\n",
+            "Step 11/32 | Loss: 0.3750 | Mean Pearson: 0.0121 | LR: 2.13e-09 | Tokens: 45,056\n",
+            "Step 13/32 | Loss: 0.4380 | Mean Pearson: -0.0126 | LR: 2.02e-09 | Tokens: 53,248\n",
             "\n",
             "Running validation at step 12...\n",
-            "  Validation Loss: 0.5107\n",
-            "  Validation Mean Pearson: -0.0596\n",
-            "    ENCFF884LDL/pearson: -0.0596\n",
-            "Step 14/32 | Loss: 0.2984 | Mean Pearson: -0.0820 | LR: 1.91e-09 | Tokens: 61,440\n",
-            "Step 16/32 | Loss: 0.5219 | Mean Pearson: -0.0668 | LR: 1.80e-09 | Tokens: 69,632\n",
+            "  Validation Loss: 0.3997\n",
+            "  Validation Mean Pearson: 0.0093\n",
+            "    ENCFF884LDL/pearson: 0.0093\n",
+            "Step 15/32 | Loss: 0.3469 | Mean Pearson: -0.0279 | LR: 1.91e-09 | Tokens: 61,440\n",
+            "Step 17/32 | Loss: 0.5098 | Mean Pearson: -0.2044 | LR: 1.80e-09 | Tokens: 69,632\n",
             "\n",
             "Running validation at step 16...\n",
-            "  Validation Loss: 0.8410\n",
-            "  Validation Mean Pearson: 0.0041\n",
-            "    ENCFF884LDL/pearson: 0.0041\n",
-            "Step 18/32 | Loss: 0.3663 | Mean Pearson: 0.0888 | LR: 1.67e-09 | Tokens: 77,824\n",
-            "Step 20/32 | Loss: 0.4024 | Mean Pearson: -0.0628 | LR: 1.54e-09 | Tokens: 86,016\n",
+            "  Validation Loss: 0.3752\n",
+            "  Validation Mean Pearson: -0.0178\n",
+            "    ENCFF884LDL/pearson: -0.0178\n",
+            "Step 19/32 | Loss: 0.4899 | Mean Pearson: -0.0424 | LR: 1.67e-09 | Tokens: 77,824\n",
+            "Step 21/32 | Loss: 0.3889 | Mean Pearson: -0.0332 | LR: 1.54e-09 | Tokens: 86,016\n",
             "\n",
             "Running validation at step 20...\n",
-            "  Validation Loss: 0.4043\n",
-            "  Validation Mean Pearson: -0.1108\n",
-            "    ENCFF884LDL/pearson: -0.1108\n",
-            "Step 22/32 | Loss: 0.4096 | Mean Pearson: -0.0249 | LR: 1.39e-09 | Tokens: 94,208\n",
-            "Step 24/32 | Loss: 0.3930 | Mean Pearson: -0.0779 | LR: 1.23e-09 | Tokens: 102,400\n",
+            "  Validation Loss: 0.4217\n",
+            "  Validation Mean Pearson: -0.0205\n",
+            "    ENCFF884LDL/pearson: -0.0205\n",
+            "Step 23/32 | Loss: 0.3392 | Mean Pearson: 0.0235 | LR: 1.39e-09 | Tokens: 94,208\n",
+            "Step 25/32 | Loss: 0.4165 | Mean Pearson: 0.0033 | LR: 1.23e-09 | Tokens: 102,400\n",
             "\n",
             "Running validation at step 24...\n",
-            "  Validation Loss: 0.3426\n",
-            "  Validation Mean Pearson: 0.0236\n",
-            "    ENCFF884LDL/pearson: 0.0236\n",
-            "Step 26/32 | Loss: 0.4457 | Mean Pearson: -0.0219 | LR: 1.04e-09 | Tokens: 110,592\n",
-            "Step 28/32 | Loss: 0.4520 | Mean Pearson: 0.0580 | LR: 8.04e-10 | Tokens: 118,784\n",
+            "  Validation Loss: 0.4363\n",
+            "  Validation Mean Pearson: -0.0379\n",
+            "    ENCFF884LDL/pearson: -0.0379\n",
+            "Step 27/32 | Loss: 0.7630 | Mean Pearson: 0.0683 | LR: 1.04e-09 | Tokens: 110,592\n",
+            "Step 29/32 | Loss: 0.7357 | Mean Pearson: 0.0050 | LR: 8.04e-10 | Tokens: 118,784\n",
             "\n",
             "Running validation at step 28...\n",
-            "  Validation Loss: 0.3757\n",
-            "  Validation Mean Pearson: 0.0050\n",
-            "    ENCFF884LDL/pearson: 0.0050\n",
-            "Step 30/32 | Loss: 0.9341 | Mean Pearson: -0.0122 | LR: 4.64e-10 | Tokens: 126,976\n",
+            "  Validation Loss: 0.6629\n",
+            "  Validation Mean Pearson: -0.0370\n",
+            "    ENCFF884LDL/pearson: -0.0370\n",
+            "Step 31/32 | Loss: 0.3690 | Mean Pearson: -0.0808 | LR: 4.64e-10 | Tokens: 126,976\n",
             "\n",
             "Training completed after 32 steps!\n"
           ]
@@ -1290,7 +1224,7 @@
         "        current_lr = scheduler.get_last_lr()[0] if scheduler else config[\"learning_rate\"]\n",
         "        print(f\"Step {optimizer_step_idx + 1}/{num_steps_training} | \"\n",
         "              f\"Loss: {avg_loss:.4f} | \"\n",
-        "              f\"Mean Pearson: {train_metrics_dict['mean/pearson']:.4f} | \"\n",
+        "              f\"Mean Pearson: {train_metrics_dict['metrics_scaled/mean/pearson']:.4f} | \"\n",
         "              f\"LR: {current_lr:.2e} | \"\n",
         "              f\"Tokens: {num_tokens_seen:,}\")\n",
         "        train_metrics.reset()\n",
@@ -1311,9 +1245,9 @@
         "        # Print validation metrics\n",
         "        val_metrics_dict = val_metrics.compute()\n",
         "        print(f\"  Validation Loss: {np.mean(val_losses):.4f}\")\n",
-        "        print(f\"  Validation Mean Pearson: {val_metrics_dict['mean/pearson']:.4f}\")\n",
+        "        print(f\"  Validation Mean Pearson: {val_metrics_dict['metrics_scaled/mean/pearson']:.4f}\")\n",
         "        for track_name in config[\"bigwig_file_ids\"]:\n",
-        "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'{track_name}/pearson']:.4f}\")\n",
+        "            print(f\"    {track_name}/pearson: {val_metrics_dict[f'metrics_scaled/{track_name}/pearson']:.4f}\")\n",
         "        \n",
         "        model.train()  # Back to training mode\n",
         "\n",
@@ -1329,7 +1263,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 44,
       "metadata": {},
       "outputs": [],
       "source": [
@@ -1381,18 +1315,36 @@
             "\n",
             "==================================================\n",
             "Test Set Evaluation\n",
-            "==================================================\n"
+            "==================================================\n",
+            "Running test evaluation with 5 steps (10 samples)\n"
           ]
         },
         {
-          "ename": "NameError",
-          "evalue": "name 'test_dataset' is not defined",
-          "output_type": "error",
-          "traceback": [
-            "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
-            "\u001b[31mNameError\u001b[39m                                 Traceback (most recent call last)",
-            "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[68]\u001b[39m\u001b[32m, line 10\u001b[39m\n\u001b[32m      8\u001b[39m \u001b[38;5;66;03m# Calculate number of test steps (based on deepspeed pipeline)\u001b[39;00m\n\u001b[32m      9\u001b[39m test_batch_size = config[\u001b[33m\"\u001b[39m\u001b[33mbatch_size\u001b[39m\u001b[33m\"\u001b[39m]\n\u001b[32m---> \u001b[39m\u001b[32m10\u001b[39m num_test_samples = \u001b[38;5;28mlen\u001b[39m(\u001b[43mtest_dataset\u001b[49m)\n\u001b[32m     11\u001b[39m num_test_steps = num_test_samples // test_batch_size\n\u001b[32m     13\u001b[39m \u001b[38;5;28mprint\u001b[39m(\u001b[33mf\u001b[39m\u001b[33m\"\u001b[39m\u001b[33mRunning test evaluation with \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnum_test_steps\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m steps (\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mnum_test_samples\u001b[38;5;132;01m}\u001b[39;00m\u001b[33m samples)\u001b[39m\u001b[33m\"\u001b[39m)\n",
-            "\u001b[31mNameError\u001b[39m: name 'test_dataset' is not defined"
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/home/y-bornachot/venvs/ntv3-env/lib/python3.12/site-packages/torch/amp/autocast_mode.py:287: UserWarning: In CPU autocast, but the target dtype is not supported. Disabling autocast.\n",
+            "CPU Autocast only supports dtype of torch.bfloat16, torch.float16 currently.\n",
+            "  warnings.warn(error_message)\n"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "==================================================\n",
+            "Test Set Results\n",
+            "==================================================\n",
+            "\n",
+            "Scaled Metrics (scaled predictions vs scaled targets):\n",
+            "  Mean Pearson (scaled): -0.0362\n",
+            "    metrics_scaled/ENCFF884LDL/pearson: -0.0362\n",
+            "\n",
+            "Raw Metrics (raw predictions vs raw targets):\n",
+            "  Mean Pearson (raw): -0.0362\n",
+            "    metrics_raw/ENCFF884LDL/pearson: -0.0362\n",
+            "==================================================\n"
           ]
         }
       ],
@@ -1432,14 +1384,14 @@
         "print(\"Test Set Results\")\n",
         "print(\"=\"*50)\n",
         "print(f\"\\nScaled Metrics (scaled predictions vs scaled targets):\")\n",
-        "print(f\"  Mean Pearson (scaled): {test_metrics_dict['mean/pearson_scaled']:.4f}\")\n",
+        "print(f\"  Mean Pearson (scaled): {test_metrics_dict['metrics_scaled/mean/pearson']:.4f}\")\n",
         "for track_name in config[\"bigwig_file_ids\"]:\n",
-        "    print(f\"    {track_name}/pearson_scaled: {test_metrics_dict[f'{track_name}/pearson_scaled']:.4f}\")\n",
+        "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'metrics_scaled/{track_name}/pearson']:.4f}\")\n",
         "\n",
         "print(f\"\\nRaw Metrics (raw predictions vs raw targets):\")\n",
-        "print(f\"  Mean Pearson (raw): {test_metrics_dict['mean/pearson_raw']:.4f}\")\n",
+        "print(f\"  Mean Pearson (raw): {test_metrics_dict['metrics_raw/mean/pearson']:.4f}\")\n",
         "for track_name in config[\"bigwig_file_ids\"]:\n",
-        "    print(f\"    {track_name}/pearson_raw: {test_metrics_dict[f'{track_name}/pearson_raw']:.4f}\")\n",
+        "    print(f\"    {track_name}/pearson: {test_metrics_dict[f'metrics_raw/{track_name}/pearson']:.4f}\")\n",
         "print(\"=\"*50)"
       ]
     },