diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
new file mode 100644
index 000000000..214c5116d
--- /dev/null
+++ b/.github/pull_request_template.md
@@ -0,0 +1,8 @@
+# Description
+<!-- Add Notes if relevant -->
+
+# Overview of Changes
+
+# Testing Notes 
+- [ ] {You should have tested the code in some way (doesn't have to be a unit test)}
+- [ ] {Describe any additional tests}
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 9993737db..6e802ed70 100644
--- a/.gitignore
+++ b/.gitignore
@@ -138,5 +138,8 @@ data/physionet.org/
 # VSCode settings
 .vscode/
 
+# Local Data
+local_data/
+
 # Model weight files (large binaries, distributed separately)
 weightfiles/
\ No newline at end of file
diff --git a/cc_state.sh b/cc_state.sh
new file mode 100644
index 000000000..61b027d87
--- /dev/null
+++ b/cc_state.sh
@@ -0,0 +1,272 @@
+#!/usr/bin/env bash
+# cc_state.sh — one-shot local script for Campus Cluster Table 2 management.
+#
+# Usage:
+#   bash cc_state.sh                 # show queue + recent job status
+#   bash cc_state.sh resubmit        # sync, cancel all, clean cache, warm, submit all 18
+#   bash cc_state.sh cancel          # cancel all pending/running jobs
+#   bash cc_state.sh sync            # sync scripts to CC only (no submit)
+#   bash cc_state.sh clean-cache     # delete corrupted parquet cache entries
+#   bash cc_state.sh results         # print best AUROC/AUPRC/F1 per completed run
+#   bash cc_state.sh logs [MODEL]    # tail recent logs (optional model filter)
+set -euo pipefail
+
+CC="${CC:-rianatri@cc-login.campuscluster.illinois.edu}"
+REMOTE_REPO="${REMOTE_REPO:-/u/rianatri/PyHealth}"
+LOCAL_REPO="${LOCAL_REPO:-$(cd "$(dirname "$0")" && pwd)}"
+SSH_KEY="${SSH_KEY:-}"   # e.g. SSH_KEY=~/.ssh/id_ed25519 bash cc_state.sh
+SSH_OPTS="-o StrictHostKeyChecking=no${SSH_KEY:+ -o IdentitiesOnly=yes -i ${SSH_KEY}}"
+CMD="${1:-state}"
+
+ssh_cc() { ssh ${SSH_OPTS} "${CC}" "$@"; }
+rsync_cc() {
+    rsync -avz --relative \
+        -e "ssh ${SSH_OPTS}" \
+        "$@" \
+        "${CC}:${REMOTE_REPO}/"
+}
+
+# ── sync ──────────────────────────────────────────────────────────────────────
+do_sync() {
+    echo "[sync] Syncing scripts and pyhealth source to CC..."
+    cd "${LOCAL_REPO}"
+    rsync_cc \
+        pyhealth/ \
+        examples/mortality_prediction/unified_embedding_e2e_mimic4.py \
+        scripts/slurm/run_table2.sh \
+        scripts/slurm/run_cachewarm.sh \
+        scripts/slurm/submit_table2_random.sh \
+        scripts/slurm/submit_table2_ic.sh \
+        scripts/slurm/setup_cc.sh \
+        scripts/condor/warm_table2_cache.py
+    echo "[sync] Done."
+}
+
+# ── state ─────────────────────────────────────────────────────────────────────
+do_state() {
+    echo "=== Queue (rianatri) ==="
+    ssh_cc "squeue -u rianatri --format='%.10i %.12P %.22j %.8T %.10M %.6D %R' 2>/dev/null || true"
+
+    echo ""
+    echo "=== Recent job outcomes (last 24h) ==="
+    ssh_cc "sacct -u rianatri --starttime=now-24hours \
+        --format=JobID%15,JobName%25,State%12,ExitCode,Elapsed \
+        --noheader 2>/dev/null | grep -v '\.batch\|\.extern' || true"
+
+    echo ""
+    echo "=== GPU availability ==="
+    ssh_cc "sinfo -p eng-research-gpu,IllinoisComputes-GPU \
+        -o '%.20P %.10T %.6D %.15G' 2>/dev/null || true"
+}
+
+# ── cancel ────────────────────────────────────────────────────────────────────
+do_cancel() {
+    echo "[cancel] Cancelling all jobs for rianatri..."
+    ssh_cc "scancel -u rianatri 2>/dev/null || true; echo '  Done.'"
+}
+
+# ── clean-cache ───────────────────────────────────────────────────────────────
+do_clean_cache() {
+    echo "[clean-cache] Removing corrupted parquet cache entries on CC..."
+    ssh_cc "bash -s" <<'EOF'
+CACHE_DIR="/u/${USER}/pyhealth_cache"
+echo "  Cache dir: ${CACHE_DIR}"
+
+# Remove any global_event_df.parquet directories that are empty or have 0-byte files
+# (these are left behind by failed dask writes)
+find "${CACHE_DIR}" -name "global_event_df.parquet" -type d | while read -r d; do
+    # Check for empty or 0-byte parquet files inside
+    bad=$(find "${d}" -name "*.parquet" -size 0 2>/dev/null | head -1)
+    if [[ -n "${bad}" ]] || [[ -z "$(ls -A "${d}" 2>/dev/null)" ]]; then
+        echo "  Removing corrupted: ${d}"
+        rm -rf "${d}"
+    else
+        echo "  OK (non-empty): ${d}"
+    fi
+done
+
+# Also clean up any stale dask temp dirs
+rm -rf /u/${USER}/dask_tmp/ 2>/dev/null && echo "  Cleaned dask_tmp" || true
+mkdir -p /u/${USER}/dask_tmp
+echo "  Done."
+EOF
+}
+
+# ── resubmit ──────────────────────────────────────────────────────────────────
+do_resubmit() {
+    do_sync
+
+    echo ""
+    echo "[resubmit] Cancelling all jobs..."
+    ssh_cc "scancel -u rianatri 2>/dev/null || true; sleep 2"
+
+    echo ""
+    do_clean_cache
+
+    echo ""
+    echo "[resubmit] Submitting cachewarm + 18 training jobs (training deps on cachewarm)..."
+    ssh_cc "REMOTE_REPO='${REMOTE_REPO}' bash -s" <<'EOF'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+
+mkdir -p /u/${USER}/dask_tmp logs/slurm
+
+# Submit cachewarm job
+WARM_JOB=$(sbatch \
+    --account=jimeng-cs-eng \
+    --partition=eng-research-gpu \
+    --nodes=1 --ntasks=1 --cpus-per-task=4 \
+    --mem=48G --gres=gpu:1 --time=08:00:00 \
+    --job-name=table2_cachewarm \
+    --output=logs/slurm/table2_cachewarm_%j.out \
+    --error=logs/slurm/table2_cachewarm_%j.err \
+    scripts/slurm/run_cachewarm.sh | awk '{print $NF}')
+echo "  Cachewarm job: ${WARM_JOB}"
+
+# All 18 jobs → IllinoisComputes-GPU (IC), chained after cachewarm.
+# mlp/rnn: BERT encoder OOMs on A10 24GB at bs=16; IC A100/H200 handles bs=16 fine.
+DEPEND="--dependency=afterok:${WARM_JOB}"
+
+for model in mlp rnn ehrmamba transformer bottleneck_transformer jambaehr; do
+    while IFS= read -r seed; do
+        case "${model}" in
+            mlp)                   bs_var="TABLE2_BS_MLP=16"      ; tl="6:00:00"  ;;
+            rnn)                   bs_var="TABLE2_BS_RNN=16"      ; tl="6:00:00"  ;;
+            ehrmamba)              bs_var="TABLE2_BS_EHRMAMBA=8"  ; tl="12:00:00" ;;
+            transformer)          bs_var="TABLE2_BS_TRANSFORMER=4"; tl="18:00:00" ;;
+            bottleneck_transformer) bs_var="TABLE2_BS_BOTTLENECK=4"; tl="18:00:00" ;;
+            jambaehr)             bs_var="TABLE2_BS_JAMBAEHR=4"  ; tl="18:00:00" ;;
+        esac
+        job=$(sbatch \
+            --job-name="t2ic_${model}_s${seed}" \
+            --account=jimeng-ic \
+            --partition=IllinoisComputes-GPU \
+            --nodes=1 --ntasks=1 --cpus-per-task=4 \
+            --mem=32G --gres=gpu:1 --time="${tl}" \
+            --output="logs/slurm/table2ic_${model}_seed${seed}_%j.out" \
+            --error="logs/slurm/table2ic_${model}_seed${seed}_%j.err" \
+            --export="ALL,MODEL=${model},SEED=${seed},${bs_var}" \
+            ${DEPEND} \
+            scripts/slurm/run_table2.sh | awk '{print $NF}')
+        echo "  Submitted ${model} seed=${seed} → ${job}"
+    done < scripts/slurm/table2_random_seeds.txt
+done
+
+echo ""
+echo "19 jobs queued (1 cachewarm + 18 training). Queue:"
+squeue -u rianatri --format="%.10i %.12P %.22j %.8T %.10M %.6D %R"
+EOF
+}
+
+# ── results ───────────────────────────────────────────────────────────────────
+do_results() {
+    echo "[results] Fetching completed results from CC..."
+    ssh_cc "REMOTE_REPO='${REMOTE_REPO}' bash -s" <<'EOSSH'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+OUT="output/table2"
+
+if [[ ! -d "${OUT}" ]]; then
+    echo "  No output directory found."
+    exit 0
+fi
+
+# Print header
+printf "\n%-35s %8s %8s %8s %8s %s\n" "Run" "AUROC" "AUPRC" "F1" "Acc" "Epochs"
+printf '%s\n' "$(printf '%.0s-' {1..80})"
+
+found=0
+for d in "${OUT}"/*/; do
+    run=$(basename "${d}")
+    json="${d}metrics_history.json"
+    [[ -f "${json}" ]] || continue
+    found=1
+    # Extract best val roc_auc (max), and corresponding auprc/f1/acc from that epoch
+    python3 - "${json}" "${run}" <<'PY'
+import json, sys
+path, run = sys.argv[1], sys.argv[2]
+with open(path) as f:
+    h = json.load(f)
+
+# Handle two possible formats:
+# Format A: {"val": [{"roc_auc": 0.8, "epoch": 1, ...}, ...]}
+# Format B: [{"epoch": 1, "roc_auc": 0.8, ...}, ...]  (flat list)
+# Format C: {"roc_auc": [0.8, 0.9, ...], "pr_auc": [...]}  (dict of lists)
+if isinstance(h, list):
+    val = h
+elif isinstance(h, dict):
+    if "val" in h and isinstance(h["val"], list) and h["val"] and isinstance(h["val"][0], dict):
+        val = h["val"]
+    elif "roc_auc" in h and isinstance(h["roc_auc"], list):
+        # dict-of-lists format
+        keys = list(h.keys())
+        n = len(h[keys[0]])
+        val = [{k: h[k][i] for k in keys} for i in range(n)]
+    else:
+        val = []
+else:
+    val = []
+
+if not val:
+    print(f"{'  '+run:<35} {'N/A':>8} {'N/A':>8} {'N/A':>8} {'N/A':>8} -")
+    sys.exit(0)
+
+def safe_get(d, *keys):
+    for k in keys:
+        if isinstance(d, dict) and k in d:
+            return d[k]
+    return 0
+
+best = max(val, key=lambda e: safe_get(e, "roc_auc"))
+epoch = safe_get(best, "epoch")
+total = len(val)
+print(f"{run:<35} {safe_get(best,'roc_auc'):8.4f} {safe_get(best,'pr_auc','auprc'):8.4f} {safe_get(best,'f1'):8.4f} {safe_get(best,'accuracy'):8.4f} {epoch}/{total}")
+PY
+done
+
+if [[ "${found}" -eq 0 ]]; then
+    echo "  No completed results yet (metrics_history.json not found in any run dir)."
+fi
+EOSSH
+}
+
+# ── logs ──────────────────────────────────────────────────────────────────────
+do_logs() {
+    local filter="${2:-}"
+    echo "[logs] Fetching recent log tails from CC (filter: ${filter:-all})..."
+    ssh_cc "REMOTE_REPO='${REMOTE_REPO}' FILTER='${filter}' bash -s" <<'EOSSH'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+LOG_DIR="logs/slurm"
+if [[ -n "${FILTER}" ]]; then
+    mapfile -t LOGS < <(ls -t "${LOG_DIR}"/*"${FILTER}"*.out 2>/dev/null | head -6)
+else
+    mapfile -t LOGS < <(ls -t "${LOG_DIR}"/*.out 2>/dev/null | head -9)
+fi
+if [[ "${#LOGS[@]}" -eq 0 ]]; then
+    echo "  No log files found."
+else
+    for f in "${LOGS[@]}"; do
+        echo ""
+        echo "━━━ ${f} ━━━"
+        tail -20 "${f}" 2>/dev/null || echo "  (empty)"
+    done
+fi
+EOSSH
+}
+
+# ── dispatch ──────────────────────────────────────────────────────────────────
+case "${CMD}" in
+    state)        do_state ;;
+    sync)         do_sync ;;
+    cancel)       do_cancel ;;
+    clean-cache)  do_clean_cache ;;
+    resubmit)     do_resubmit ;;
+    results)      do_results ;;
+    logs)         do_logs "$@" ;;
+    *)
+        echo "Unknown command: ${CMD}"
+        echo "Usage: bash cc_state.sh [state|sync|cancel|clean-cache|resubmit|results|logs [MODEL]]"
+        exit 1
+        ;;
+esac
diff --git a/clustersetup.md b/clustersetup.md
new file mode 100644
index 000000000..d9fdc1563
--- /dev/null
+++ b/clustersetup.md
@@ -0,0 +1,42 @@
+# Running PyHealth on the Campus Cluster
+
+## Quick Start
+
+```bash
+git clone https://github.com/Multimodal-PyHealth/PyHealth.git
+cd PyHealth
+chmod +x setup.sh
+./setup.sh
+```
+
+## Data Paths
+
+All MIMIC-4 data is under `/projects/illinois/eng/cs/jimeng/physionet.org/files/`:
+
+| Data | Path |
+| EHR | `/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2` |
+| Clinical Notes | `/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note` |
+| Chest X-rays | `/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-cxr-jpg/2.1.0` |
+
+**Important:** `NOTE_ROOT` should be `.../mimic-note` (not `.../mimic-note/note`). The config YAML appends `note/` automatically.
+
+Set `CACHE_DIR` to your own writable directory: `/u/<NetID>/pyhealth_cache`
+
+## Running on a Compute Node
+
+Run on a compute node:
+Slurm command example:
+```bash
+srun --account=jimeng-cs-eng --partition=eng-research-gpu --time=00:10:00 --gres=gpu:1 --pty bash
+```
+
+Once on the compute node, re-activate and run:
+
+```bash
+module load miniconda3/24.9.2
+conda activate pyhealth2
+cd ~/PyHealth
+python examples/mortality_prediction/multimodal_mimic4.py
+```
+For a clean install, which reruns the setup proccess by deleting the enviroment packages and conda env and reinstalls them, run:
+```./setup.sh clean```
\ No newline at end of file
diff --git a/condor_setup.sh b/condor_setup.sh
new file mode 100755
index 000000000..f76c05a1a
--- /dev/null
+++ b/condor_setup.sh
@@ -0,0 +1,108 @@
+#!/bin/bash
+set -e #sets flag so script stops on any error
+
+ENV_NAME="pyhealth2"
+
+resolve_conda_sh() {
+    # 1) explicit override
+    if [ -n "${CONDA_SH:-}" ] && [ -f "${CONDA_SH}" ]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+
+    # 2) conda on PATH
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [ -n "${base}" ] && [ -f "${base}/etc/profile.d/conda.sh" ]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+
+    # 3) optional module systems (quiet)
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [ -n "${mod_base}" ] && [ -f "${mod_base}/etc/profile.d/conda.sh" ]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+
+    # 4) common install locations
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [ -f "${c}" ]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+
+    # 5) broad filesystem fallback for unknown install layouts
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [ -n "${found}" ] && [ -f "${found}" ]; then
+        echo "${found}"
+        return 0
+    fi
+
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [ -z "${CONDA_SH}" ] || [ ! -f "${CONDA_SH}" ]; then
+    echo "ERROR: conda.sh not found." >&2
+    echo "Set it explicitly, e.g.: export CONDA_SH=/path/to/conda.sh" >&2
+    exit 1
+fi
+
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+
+# fresh install if requested
+if [ "${1:-}" == "clean" ]; then
+    conda deactivate 2>/dev/null || true
+    conda env remove -n ${ENV_NAME} -y 2>/dev/null || true
+    echo "Removed env ${ENV_NAME}"
+fi
+
+# Create env if missing
+if conda env list | awk '{print $1}' | grep -Fxq "${ENV_NAME}"; then
+    echo "Environment '${ENV_NAME}' already exists."
+else
+    conda create -n ${ENV_NAME} python=3.12 -y
+fi
+
+conda activate ${ENV_NAME}
+
+# Ensure required runtime deps are present even if env pre-existed.
+if ! python -c "import torch" >/dev/null 2>&1; then
+    pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124
+fi
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    pip install -e . #installs the current directory as an editable package, a pointer to the local source code
+fi
+
+# Verify
+python -c "import pyhealth; print('PyHealth: OK')"
+python -c "import torch; print(f'PyTorch {torch.__version__}, CUDA: {torch.cuda.is_available()}')"
diff --git a/configs/train/base.yaml b/configs/train/base.yaml
new file mode 100644
index 000000000..70c7f2c7c
--- /dev/null
+++ b/configs/train/base.yaml
@@ -0,0 +1,80 @@
+# PyHealth unified training config — base defaults
+# All values here are the OOM-safe defaults validated on devsplit.
+# Override per-condition and per-model in condition/model-specific configs.
+
+# ── paths ──────────────────────────────────────────────────────────────────────
+ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2
+note_root: /shared/rsaas/physionet.org/files/mimic-note
+cache_dir: /home/rianatri/pyhealth_cache
+output_dir: output/unified
+
+# ── task ──────────────────────────────────────────────────────────────────────
+# One of: notes_labs | labs_only | icd_labs | clinical_notes_icd_labs | stagenet
+task: notes_labs
+observation_window_hours: 24
+
+# Task flags (notes_labs only)
+icd_codes: false
+include_vitals: false
+balanced_sampling: false
+balanced_ratio: 1.0
+
+# ── model ─────────────────────────────────────────────────────────────────────
+# One of: mlp | rnn | transformer | bottleneck_transformer | ehrmamba | jambaehr
+model: mlp
+freeze_encoder: false
+
+# Shared embedding dims — safe defaults across all models on 24 GB GPU (frozen)
+# or 80 GB A100 (full BERT).
+embedding_dim: 128
+hidden_dim: 128
+heads: 4
+num_layers: 2
+dropout: 0.1
+
+# ── training ──────────────────────────────────────────────────────────────────
+epochs: 50
+batch_size: 16
+lr: null          # null = model-specific default (1e-4 for all)
+weight_decay: 1.0e-5
+patience: 10
+seed: 42
+num_workers: 2
+
+# ── dev mode ──────────────────────────────────────────────────────────────────
+# 0 = full dataset; N > 0 = limit to N patients (devsplit)
+dev: 0
+
+# ── model-specific overrides ──────────────────────────────────────────────────
+# These are merged at runtime based on `model` value.
+# Keys match CLI args without the leading `--`.
+_model_overrides:
+  transformer:
+    batch_size: 2
+    embedding_dim: 64
+    hidden_dim: 64
+    heads: 2
+    num_layers: 1
+  bottleneck_transformer:
+    batch_size: 2
+    embedding_dim: 96
+    hidden_dim: 96
+    heads: 2
+    num_layers: 1
+    max_grad_norm: 0.5
+    bottlenecks_n: 4
+    fusion_startidx: 1
+  ehrmamba:
+    batch_size: 2
+    embedding_dim: 96
+    hidden_dim: 96
+    mamba_state_size: 16
+    mamba_conv_kernel: 4
+  jambaehr:
+    batch_size: 2
+    embedding_dim: 64
+    hidden_dim: 64
+    jamba_transformer_layers: 1
+    jamba_mamba_layers: 2
+    mamba_state_size: 16
+    mamba_conv_kernel: 4
diff --git a/configs/train/e2e_balanced.yaml b/configs/train/e2e_balanced.yaml
new file mode 100644
index 000000000..74aaedcb9
--- /dev/null
+++ b/configs/train/e2e_balanced.yaml
@@ -0,0 +1,8 @@
+# PyHealth E2E balanced condition — notes+labs with 1:1 pos:neg undersampling
+_inherit: base.yaml
+
+task: notes_labs
+balanced_sampling: true
+balanced_ratio: 1.0
+epochs: 50
+output_dir: output/e2e_full/balanced
diff --git a/configs/train/e2e_baseline.yaml b/configs/train/e2e_baseline.yaml
new file mode 100644
index 000000000..df1a6243e
--- /dev/null
+++ b/configs/train/e2e_baseline.yaml
@@ -0,0 +1,11 @@
+# PyHealth E2E baseline condition — notes+labs, ICD off, no balancing
+# Full dataset, 50 epochs, patience=10.
+
+_inherit: base.yaml
+
+task: notes_labs
+icd_codes: false
+include_vitals: false
+balanced_sampling: false
+epochs: 50
+output_dir: output/e2e_full/baseline
diff --git a/configs/train/e2e_icd_on.yaml b/configs/train/e2e_icd_on.yaml
new file mode 100644
index 000000000..13f817fd3
--- /dev/null
+++ b/configs/train/e2e_icd_on.yaml
@@ -0,0 +1,7 @@
+# PyHealth E2E ICD-on condition — notes+labs+ICD (discharge-coded leakage ablation)
+_inherit: base.yaml
+
+task: notes_labs
+icd_codes: true
+epochs: 50
+output_dir: output/e2e_full/icd_on
diff --git a/configs/train/e2e_labs_only.yaml b/configs/train/e2e_labs_only.yaml
new file mode 100644
index 000000000..d0ca5ade8
--- /dev/null
+++ b/configs/train/e2e_labs_only.yaml
@@ -0,0 +1,6 @@
+# PyHealth E2E labs-only condition — EHR-only reference baseline
+_inherit: base.yaml
+
+task: labs_only
+epochs: 50
+output_dir: output/e2e_full/labs_only
diff --git a/configs/train/smoke.yaml b/configs/train/smoke.yaml
new file mode 100644
index 000000000..af0079770
--- /dev/null
+++ b/configs/train/smoke.yaml
@@ -0,0 +1,14 @@
+# PyHealth smoke test config — devsplit, fast iteration
+# Inherits base.yaml defaults; overrides for fast validation.
+
+_inherit: base.yaml
+
+task: notes_labs
+model: mlp
+dev: 1000
+epochs: 5
+patience: 5
+batch_size: 16
+embedding_dim: 64
+hidden_dim: 64
+output_dir: output/smoke
diff --git a/docs/api/models.rst b/docs/api/models.rst
index 7368dec94..568e5cfb6 100644
--- a/docs/api/models.rst
+++ b/docs/api/models.rst
@@ -176,6 +176,7 @@ API Reference
     models/pyhealth.models.RNN
     models/pyhealth.models.GNN
     models/pyhealth.models.Transformer
+    models/pyhealth.models.BottleneckTransformer
     models/pyhealth.models.TransformersModel
     models/pyhealth.models.RETAIN
     models/pyhealth.models.GAMENet
diff --git a/docs/api/tasks.rst b/docs/api/tasks.rst
index 399b8f1aa..347ec4df3 100644
--- a/docs/api/tasks.rst
+++ b/docs/api/tasks.rst
@@ -229,3 +229,4 @@ Available Tasks
     Mutation Pathogenicity (COSMIC) <tasks/pyhealth.tasks.MutationPathogenicityPrediction>
     Cancer Survival Prediction (TCGA) <tasks/pyhealth.tasks.CancerSurvivalPrediction>
     Cancer Mutation Burden (TCGA) <tasks/pyhealth.tasks.CancerMutationBurden>
+    Multimodal Mortality Prediction (MIMIC-IV) <tasks/pyhealth.tasks.multimodal_mimic4>
diff --git a/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst b/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst
new file mode 100644
index 000000000..b501e3f4e
--- /dev/null
+++ b/docs/api/tasks/pyhealth.tasks.multimodal_mimic4.rst
@@ -0,0 +1,27 @@
+pyhealth.tasks.multimodal_mimic4
+===================================
+
+.. autoclass:: pyhealth.tasks.multimodal_mimic4.ICDLabsMIMIC4
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.tasks.multimodal_mimic4.NotesLabsMIMIC4
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.tasks.multimodal_mimic4.NotesLabsCXRMIMIC4
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.tasks.multimodal_mimic4.LabsMIMIC4
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+.. autoclass:: pyhealth.tasks.multimodal_mimic4.CXRMIMIC4
+    :members:
+    :undoc-members:
+    :show-inheritance:
diff --git a/examples/bottleneck_transformer_tutorial.ipynb b/examples/bottleneck_transformer_tutorial.ipynb
new file mode 100644
index 000000000..2885be943
--- /dev/null
+++ b/examples/bottleneck_transformer_tutorial.ipynb
@@ -0,0 +1,199 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "intro_md",
+   "metadata": {},
+   "source": [
+    "# Bottleneck Transformer Tutorial\n",
+    "\n",
+    "This notebook demonstrates how to use the `BottleneckTransformer` model for multimodal healthcare data fusion in PyHealth.\n",
+    "\n",
+    "**Overview:**\n",
+    "- Initialize BottleneckTransformer with multi-modality data\n",
+    "- Demonstrate modality-specific pre-fusion vs multimodal bottleneck fusion\n",
+    "- Highlight architecture hyperparameters `bottlenecks_n` and `fusion_startidx`\n",
+    "- Inspect forward passes and probability mappings"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "env_md",
+   "metadata": {},
+   "source": [
+    "## 1. Environment Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "env_code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\")\n",
+    "\n",
+    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
+    "print(f\"Running on device: {device}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "data_md",
+   "metadata": {},
+   "source": [
+    "## 2. Data Preparation\n",
+    "We use PyHealth's `create_sample_dataset` to generate a lightweight multimodal dataset. You can substitute this with `MIMIC3Dataset`, `MIMIC4Dataset` or `OMOPDataset` for real-world scenarios."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "data_code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyhealth.datasets import create_sample_dataset\n",
+    "\n",
+    "samples = [\n",
+    "    {\n",
+    "        \"patient_id\": \"patient-0\",\n",
+    "        \"visit_id\": \"visit-0\",\n",
+    "        \"conditions\": [\"A\", \"B\", \"C\"],\n",
+    "        \"procedures\": [\"X\", \"Y\"],\n",
+    "        \"labs\": [1.0, 2.0, 3.0],\n",
+    "        \"label\": 1,\n",
+    "    },\n",
+    "    {\n",
+    "        \"patient_id\": \"patient-1\",\n",
+    "        \"visit_id\": \"visit-0\",\n",
+    "        \"conditions\": [\"D\", \"E\"],\n",
+    "        \"procedures\": [\"Y\"],\n",
+    "        \"labs\": [4.0, 5.0, 6.0],\n",
+    "        \"label\": 0,\n",
+    "    },\n",
+    "]\n",
+    "\n",
+    "input_schema = {\n",
+    "    \"conditions\": \"sequence\",\n",
+    "    \"procedures\": \"sequence\",\n",
+    "    \"labs\": \"tensor\",\n",
+    "}\n",
+    "output_schema = {\"label\": \"binary\"}\n",
+    "\n",
+    "dataset = create_sample_dataset(\n",
+    "    samples=samples,\n",
+    "    input_schema=input_schema,\n",
+    "    output_schema=output_schema,\n",
+    "    dataset_name=\"test\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "loader_md",
+   "metadata": {},
+   "source": [
+    "## 3. Dataloader Setup\n",
+    "We use PyHealth's automatic `get_dataloader` utility which converts the structured processed fields into batches."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "loader_code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyhealth.datasets import get_dataloader\n",
+    "\n",
+    "train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "model_md",
+   "metadata": {},
+   "source": [
+    "## 4. Initialize Bottleneck Transformer\n",
+    "The model initializes modality-specific transformer paths and limits the dense attention flow to bottleneck tokens specifically. \n",
+    "\n",
+    "- `fusion_startidx` parameter decides which layer cross-attention over bottlenecks activates. Lower means earlier fusion.\n",
+    "- `bottlenecks_n` regulates how many tokens represent the capacity of the bottleneck."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "model_code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyhealth.models import BottleneckTransformer\n",
+    "\n",
+    "model = BottleneckTransformer(\n",
+    "    dataset=dataset,\n",
+    "    embedding_dim=128,\n",
+    "    bottlenecks_n=4,\n",
+    "    fusion_startidx=1,\n",
+    "    num_layers=3,\n",
+    "    heads=4\n",
+    ").to(device)\n",
+    "\n",
+    "print(\"Model modalities:\", model.feature_keys)\n",
+    "print(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "forward_md",
+   "metadata": {},
+   "source": [
+    "## 5. Forward Pass\n",
+    "Perform a simple mapping to inspect outputs. PyHealth models produce unified dicts returning `loss`, probability spaces `y_prob`, and predictions `logit`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "forward_code",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_batch = next(iter(train_loader))\n",
+    "outputs = model(**data_batch)\n",
+    "\n",
+    "for k, v in outputs.items():\n",
+    "    try:\n",
+    "        print(f\"{k}: {v.shape}\")\n",
+    "    except AttributeError:\n",
+    "        print(f\"{k}: {v}\")\n",
+    "\n",
+    "print(\"\\nForward pass successful!\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
\ No newline at end of file
diff --git a/examples/cxr/mimic4_cxr_sunlab_stats.py b/examples/cxr/mimic4_cxr_sunlab_stats.py
new file mode 100644
index 000000000..eb70f12a9
--- /dev/null
+++ b/examples/cxr/mimic4_cxr_sunlab_stats.py
@@ -0,0 +1,80 @@
+"""Example script to load Sunlab MIMIC-CXR through MIMIC4Dataset
+and print stats.
+
+Usage:
+  python examples/cxr/mimic4_cxr_sunlab_stats.py \
+    --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \
+        --note-root /shared/rsaas/physionet.org/files/mimic-note \
+        --cxr-root /shared/rsaas/physionet.org/files/MIMIC-CXR \
+        --cache-dir /shared/eng/pyhealth
+"""
+
+import argparse
+
+from pyhealth.datasets import MIMIC4Dataset
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Load Sunlab MIMIC-CXR variant and run dataset.stats()."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+        help="Root directory for MIMIC-IV EHR data.",
+    )
+    parser.add_argument(
+        "--note-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimic-note",
+        help="Root directory for MIMIC-IV notes data.",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+        help="Root directory for the Sunlab MIMIC-CXR mirror.",
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+        help="Optional cache directory for PyHealth dataset artifacts.",
+    )
+    parser.add_argument(
+        "--num-workers",
+        type=int,
+        default=4,
+        help="Number of workers for dataset processing.",
+    )
+    parser.add_argument(
+        "--dev",
+        action="store_true",
+        help="Enable dev mode (small subset).",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+
+    dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        note_root=args.note_root,
+        cxr_root=args.cxr_root,
+        cxr_variant="sunlab",
+        ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"],
+        note_tables=["discharge", "radiology"],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        num_workers=args.num_workers,
+        dev=args.dev,
+    )
+
+    # Prints table/patient/event statistics from BaseDataset.
+    dataset.stats()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/mortality_prediction/multimodal_dataset_stats.py b/examples/mortality_prediction/multimodal_dataset_stats.py
new file mode 100644
index 000000000..cb167016e
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_dataset_stats.py
@@ -0,0 +1,405 @@
+"""Standalone multimodal dataset statistics auditor.
+
+Loads a MIMIC-IV multimodal task dataset and reports per-modality missingness
+rates and token/element counts — no model, no GPU, no trainer required.
+
+Works directly on the processed SampleDataset. Processed schema (single sample,
+no batch dim):
+  discharge_note_times / radiology_note_times:
+    (input_ids, attn_mask, token_type_ids, time, type_tag)
+    input_ids shape: (N_notes, 128)  attn_mask shape: (N_notes, 128)
+    attn_mask.sum() gives real (non-padding) tokens seen by the model.
+    Note: each note is independently truncated to 128 wordpieces — no chunking.
+  icd_codes:
+    (time, value)  value shape: (N_visits, vocab_size)  [multi-hot]
+  labs_mask:
+    (time, value)  value shape: (N_timesteps, 10)  [bool/float]
+  cxr_image_times:
+    (image, time, paths)  image shape: (N_images, 3, 224, 224)
+
+Usage:
+    python examples/mortality_prediction/multimodal_dataset_stats.py \\
+        --dev --quick-test
+
+    python examples/mortality_prediction/multimodal_dataset_stats.py \\
+        --task CXRMIMIC4 --output-csv /tmp/stats.csv
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from typing import Any, Dict, List, Tuple
+
+import numpy as np
+
+from pyhealth.datasets import MIMIC4Dataset
+from pyhealth.tasks.multimodal_mimic4 import (
+    CXRMIMIC4,
+    ICDLabsMIMIC4,
+    LabsMIMIC4,
+    NotesLabsMIMIC4,
+)
+
+TASK_MAP = {
+    "ICDLabsMIMIC4": ICDLabsMIMIC4,
+    "NotesLabsMIMIC4": NotesLabsMIMIC4,
+    "LabsMIMIC4": LabsMIMIC4,
+    "CXRMIMIC4": CXRMIMIC4,
+}
+
+# Per-task data requirements: which EHR tables, notes, and CXR the task reads.
+TASK_REQUIREMENTS = {
+    "ICDLabsMIMIC4": {
+        "ehr_tables": ["diagnoses_icd", "procedures_icd", "labevents"],
+        "notes": False,
+        "cxr": False,
+    },
+    "NotesLabsMIMIC4": {
+        "ehr_tables": ["labevents"],
+        "notes": True,
+        "cxr": False,
+    },
+    "LabsMIMIC4": {
+        "ehr_tables": ["labevents"],
+        "notes": False,
+        "cxr": False,
+    },
+    "CXRMIMIC4": {
+        "ehr_tables": [],
+        "notes": False,
+        "cxr": True,
+    },
+}
+
+# bert-base-uncased encodes "[MISSING_TEXT]" to ~7 tokens with padding.
+# Real notes are always longer. Used to detect the missingness sentinel.
+_MISSING_NOTE_TOKEN_THRESHOLD = 15
+
+
+def _note_stats(tup: tuple, note_max_len: int) -> Tuple[bool, int, int]:
+    """Stats from a processed note tuple.
+
+    Schema: (input_ids, attn_mask, token_type_ids, time, type_tag)
+      input_ids shape: (N_notes, note_max_len)
+      attn_mask shape: (N_notes, note_max_len)
+
+    Returns (is_missing, n_notes, seen_tokens_total).
+    seen_tokens = attn_mask.sum() — real non-padding tokens, already capped
+    at note_max_len by the processor's truncation.
+    """
+    input_ids, attn_mask = tup[0], tup[1]
+    n_notes = input_ids.shape[0]
+    seen = int(attn_mask.sum().item())
+    is_missing = (n_notes == 1) and (seen < _MISSING_NOTE_TOKEN_THRESHOLD)
+    if is_missing:
+        return True, 0, 0
+    return False, n_notes, seen
+
+
+def _icd_stats(tup: tuple) -> Tuple[bool, int, int]:
+    """Stats from a processed icd_codes tuple.
+
+    Schema: (time, value)  value shape: (N_visits, vocab_size) [multi-hot]
+
+    Returns (is_missing, n_visits, n_code_activations).
+    """
+    value = tup[1]
+    n_visits = value.shape[0]
+    n_act = int(value.sum().item())
+    return n_act == 0, n_visits, n_act
+
+
+def _labs_stats(tup: tuple) -> Tuple[bool, int]:
+    """Stats from a processed labs_mask tuple.
+
+    Schema: (time, value)  value shape: (N_timesteps, 10) [bool/float]
+
+    Returns (is_missing, n_observed).
+    """
+    value = tup[1]
+    n_obs = int(value.sum().item())
+    return n_obs == 0, n_obs
+
+
+def _cxr_stats(tup: tuple, patch_count: int) -> Tuple[bool, int, int]:
+    """Stats from a processed cxr_image_times tuple.
+
+    Schema: (image, time, paths)  image shape: (N_images, 3, H, W)
+
+    Returns (is_missing, n_images, cxr_tokens).
+    A zero-valued image tensor indicates the missing-image sentinel.
+    """
+    image = tup[0]
+    n_images = image.shape[0]
+    if n_images == 0 or float(image.sum()) < 1e-8:
+        return True, 0, 0
+    return False, n_images, n_images * patch_count
+
+
+def audit_sample(
+    sample: Dict[str, Any],
+    note_max_len: int,
+    cxr_patch_count: int,
+) -> Dict[str, Any]:
+    """Audit one processed SampleDataset sample dict."""
+    row: Dict[str, Any] = {}
+
+    for note_key in ("discharge_note_times", "radiology_note_times", "admission_note_times"):
+        if note_key not in sample:
+            continue
+        miss, n_notes, seen_tok = _note_stats(
+            sample[note_key], note_max_len
+        )
+        row[f"{note_key}__missing"] = int(miss)
+        row[f"{note_key}__n_notes"] = n_notes
+        row[f"{note_key}__seen_tokens"] = seen_tok
+
+    if "icd_codes" in sample:
+        miss, n_visits, n_act = _icd_stats(sample["icd_codes"])
+        row["icd_codes__missing"] = int(miss)
+        row["icd_codes__n_visits"] = n_visits
+        row["icd_codes__n_activations"] = n_act
+
+    if "labs_mask" in sample:
+        miss, n_obs = _labs_stats(sample["labs_mask"])
+        row["labs__missing"] = int(miss)
+        row["labs__n_obs"] = n_obs
+
+    if "cxr_image_times" in sample:
+        miss, n_img, cxr_tok = _cxr_stats(
+            sample["cxr_image_times"], cxr_patch_count
+        )
+        row["cxr__missing"] = int(miss)
+        row["cxr__n_images"] = n_img
+        row["cxr__tokens"] = cxr_tok
+
+    note_tok = sum(
+        row.get(f"{k}__seen_tokens", 0)
+        for k in ("discharge_note_times", "radiology_note_times", "admission_note_times")
+    )
+    row["total_tokens"] = (
+        note_tok
+        + row.get("icd_codes__n_activations", 0)
+        + row.get("labs__n_obs", 0)
+        + row.get("cxr__tokens", 0)
+    )
+    return row
+
+
+def _stats(arr: np.ndarray) -> str:
+    if len(arr) == 0:
+        return "N/A"
+    return (
+        f"mean={arr.mean():.1f}  median={np.median(arr):.0f}"
+        f"  p90={np.percentile(arr, 90):.0f}  max={arr.max():.0f}"
+    )
+
+
+def print_report(rows: List[Dict], args: argparse.Namespace) -> None:
+    n = len(rows)
+    print()
+    print(f"Task: {args.task}   Samples: {n:,}   dev={args.dev}")
+    print()
+
+    def _arr(key: str) -> np.ndarray:
+        return np.array([r[key] for r in rows if key in r], dtype=float)
+
+    hdr = (
+        f"{'Modality':<28} {'missing%':>10}"
+        f"  {'mean':>8}  {'median':>8}  {'p90':>8}  {'max':>8}"
+    )
+    print(hdr)
+    print("-" * len(hdr))
+
+    modality_specs = []
+    for note_key, label in [
+        ("discharge_note_times", "discharge notes"),
+        ("radiology_note_times", "radiology notes"),
+        ("admission_note_times", "admission notes"),
+    ]:
+        miss = _arr(f"{note_key}__missing")
+        if len(miss):
+            modality_specs.append(
+                (label, miss, _arr(f"{note_key}__n_notes"))
+            )
+
+    if rows and "icd_codes__missing" in rows[0]:
+        modality_specs.append((
+            "icd codes",
+            _arr("icd_codes__missing"),
+            _arr("icd_codes__n_activations"),
+        ))
+    if rows and "labs__missing" in rows[0]:
+        modality_specs.append((
+            "labs (observations)",
+            _arr("labs__missing"),
+            _arr("labs__n_obs"),
+        ))
+    if rows and "cxr__missing" in rows[0]:
+        modality_specs.append((
+            "cxr images",
+            _arr("cxr__missing"),
+            _arr("cxr__n_images"),
+        ))
+
+    for label, miss, counts in modality_specs:
+        miss_pct = f"{miss.mean() * 100:.1f}%"
+        print(
+            f"{label:<28} {miss_pct:>10}  "
+            f"{counts.mean():>8.1f}  {np.median(counts):>8.0f}  "
+            f"{np.percentile(counts, 90):>8.0f}  {counts.max():>8.0f}"
+        )
+
+    print()
+    print(f"Note seen-tokens (cap@{args.note_max_len}, from attn_mask):")
+    for note_key, label in [
+        ("discharge_note_times", "  discharge"),
+        ("radiology_note_times", "  radiology"),
+        ("admission_note_times", "  admission"),
+    ]:
+        seen = _arr(f"{note_key}__seen_tokens")
+        if len(seen) == 0:
+            continue
+        pct_full = (seen >= args.note_max_len).mean() * 100
+        print(
+            f"{label}: {_stats(seen)}"
+            f"   % hitting cap={pct_full:.0f}%"
+        )
+
+    if rows and "cxr__tokens" in rows[0]:
+        cxr_tok = _arr("cxr__tokens")
+        print()
+        print(f"CXR tokens (@{args.cxr_patch_count} patches/image):")
+        print(f"  {_stats(cxr_tok)}")
+
+    total = _arr("total_tokens")
+    print()
+    print("Aggregate tokens/sample:")
+    print(f"  {_stats(total)}")
+    print()
+
+
+def run(args: argparse.Namespace) -> None:
+    task_cls = TASK_MAP[args.task]
+    reqs = TASK_REQUIREMENTS[args.task]
+    needs_notes = reqs["notes"]
+    needs_cxr = reqs["cxr"]
+
+    ehr_tables = reqs["ehr_tables"]
+    note_tables = ["discharge", "radiology"] if needs_notes else []
+    cxr_tables = (
+        ["metadata", "negbio", "chexpert", "split"] if needs_cxr else []
+    )
+
+    print("Loading MIMIC4Dataset ...")
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        note_root=args.note_root if needs_notes else None,
+        cxr_root=args.cxr_root if needs_cxr else None,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=ehr_tables,
+        note_tables=note_tables,
+        cxr_tables=cxr_tables,
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = task_cls(window_hours=args.observation_window_hours)
+    print("Running set_task ...")
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+    total = len(sample_dataset)
+    print(f"Total samples: {total:,}")
+
+    limit = (
+        min(args.sample_limit, total)
+        if args.sample_limit and args.sample_limit > 0
+        else total
+    )
+    print(f"Auditing {limit:,} samples ...")
+
+    rows: List[Dict] = []
+    for i in range(limit):
+        if i % 5000 == 0 and i > 0:
+            print(f"  {i:,} / {limit:,} ...")
+        rows.append(
+            audit_sample(
+                sample_dataset[i],
+                args.note_max_len,
+                args.cxr_patch_count,
+            )
+        )
+
+    if not rows:
+        print("No samples. Check roots/tables/task combination.")
+        return
+
+    print_report(rows, args)
+
+    if args.output_csv:
+        all_keys = list(rows[0].keys())
+        with open(args.output_csv, "w", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=all_keys)
+            writer.writeheader()
+            writer.writerows(rows)
+        print(f"Per-sample CSV written to: {args.output_csv}")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Audit modality missingness and token counts "
+            "for MIMIC-IV multimodal tasks."
+        )
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--note-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimic-note",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir", type=str, default="/shared/eng/pyhealth"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="NotesLabsMIMIC4",
+        choices=list(TASK_MAP.keys()),
+    )
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--note-max-len", type=int, default=128)
+    parser.add_argument("--cxr-patch-count", type=int, default=196)
+    parser.add_argument("--sample-limit", type=int, default=None)
+    parser.add_argument("--output-csv", type=str, default=None)
+
+    args = parser.parse_args()
+    if args.quick_test:
+        args.dev = True
+        if args.sample_limit is None:
+            args.sample_limit = 50
+    return args
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run(args)
diff --git a/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py
new file mode 100644
index 000000000..34616749d
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_bottleneck_mimic4_cxr.py
@@ -0,0 +1,351 @@
+"""Unified multimodal embedding + BottleneckTransformer runner.
+
+Runs chest X-ray (CXR)-only mortality prediction (MIMIC-IV + CXR) with CXRMIMIC4.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_bottleneck_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_bottleneck_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import UnifiedMultimodalEmbeddingModel
+from pyhealth.models.bottleneck_transformer import BottleneckTransformer
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(
+        dataset, [0.8, 0.1, 0.1], seed=seed
+    )
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(
+            dataset, [0.8, 0.1, 0.1], seed=seed
+        )
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_layers{args.num_layers}"
+        f"_heads{args.heads}"
+        f"_bottleneck{args.bottlenecks_n}"
+        f"_fuse{args.fusion_startidx}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(),
+        "output",
+        "multimodal_embedding_bottleneck_mimic4_cxr",
+        run_tag,
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(
+            f"  - {key}: {type(processor).__name__}, "
+            f"schema={processor.schema()}"
+        )
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = BottleneckTransformer(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        bottlenecks_n=args.bottlenecks_n,
+        fusion_startidx=args.fusion_startidx,
+        num_layers=args.num_layers,
+        heads=args.heads,
+        dropout=args.dropout,
+        unified_embedding=unified,
+    )
+    print(
+        f"BottleneckTransformer bottlenecks_n={args.bottlenecks_n}, "
+        f"fusion_startidx={args.fusion_startidx}"
+    )
+
+    train_loader = get_dataloader(
+        train_ds, batch_size=args.batch_size, shuffle=True
+    )
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: "
+        f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(
+                    f"      tuple[{i}] type={type(elem).__name__} "
+                    f"shape={shape}"
+                )
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description=(
+            "Run unified multimodal embedding + BottleneckTransformer "
+            "on MIMIC-IV mortality with CXR."
+        )
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--num-layers", type=int, default=2)
+    parser.add_argument("--heads", type=int, default=1)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--bottlenecks-n", type=int, default=4)
+    parser.add_argument("--fusion-startidx", type=int, default=1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(
+        f"Inference completed (patient_ids={num_patient_ids}, "
+        f"rows={num_rows})."
+    )
diff --git a/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py
new file mode 100644
index 000000000..f1f401096
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py
@@ -0,0 +1,330 @@
+"""Unified multimodal embedding + JambaEHR runner for MIMIC-IV + CXR.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+JambaEHR is a hybrid Transformer + Mamba architecture that interleaves
+attention and SSM (state-space model) blocks.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_jamba_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_jamba_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import JambaEHR, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_attn{args.num_transformer_layers}"
+        f"_mamba{args.num_mamba_layers}"
+        f"_heads{args.heads}"
+        f"_state{args.state_size}"
+        f"_conv{args.conv_kernel}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(), "output", "multimodal_embedding_jamba_mimic4_cxr", run_tag
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = JambaEHR(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        num_transformer_layers=args.num_transformer_layers,
+        num_mamba_layers=args.num_mamba_layers,
+        heads=args.heads,
+        dropout=args.dropout,
+        state_size=args.state_size,
+        conv_kernel=args.conv_kernel,
+        unified_embedding=unified,
+    )
+    print(f"JambaEHR unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + JambaEHR on "
+        "MIMIC-IV mortality with CXR."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--num-transformer-layers", type=int, default=2)
+    parser.add_argument("--num-mamba-layers", type=int, default=6)
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--state-size", type=int, default=16)
+    parser.add_argument("--conv-kernel", type=int, default=4)
+    parser.add_argument("--dropout", type=float, default=0.3)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py
new file mode 100644
index 000000000..23d9a9de0
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py
@@ -0,0 +1,223 @@
+"""Unified multimodal embedding + EHRMamba runner for MIMIC-IV.
+
+This script is designed to be easy to run and easy to test.
+
+Default roots are set to the local shared PhysioNet mounts:
+- ehr_root: /shared/rsaas/physionet.org/files/mimiciv/2.2
+- note_root: /shared/rsaas/physionet.org/files/mimic-note/
+
+Quick start:
+    python examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py \
+      --quick-test
+
+Full run:
+    python examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/multimodal_embedding_mamba_mimic4.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import EHRMamba, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks import NotesLabsMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:   {args.ehr_root}")
+    print(f"  note_root:  {args.note_root}")
+    print(f"  cache_dir:  {args.cache_dir}")
+    print(f"  num_workers:{args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        note_root=args.note_root,
+        ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"],
+        note_tables=["discharge", "radiology"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+    task = NotesLabsMIMIC4(
+        window_hours=args.observation_window_hours, include_icd=True
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = EHRMamba(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        num_layers=args.num_layers,
+        state_size=args.state_size,
+        conv_kernel=args.conv_kernel,
+        dropout=args.dropout,
+        unified_embedding=unified,
+    )
+    print(f"EHRMamba unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    # Debug one collated batch to verify schema-indexed fields are tensors.
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["accuracy"],
+        device=args.device,
+        enable_logging=False,
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+        )
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + EHRMamba on MIMIC-IV "
+        "mortality."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--note-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimic-note/",
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth_agent/cache",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--num-layers", type=int, default=2)
+    parser.add_argument("--state-size", type=int, default=16)
+    parser.add_argument("--conv-kernel", type=int, default=4)
+    parser.add_argument("--dropout", type=float, default=0.1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:0")
+    parser.add_argument("--num-workers", type=int, default=4)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    args = parser.parse_args()
+
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 8)
+
+    if args.device.startswith("cuda") and not torch.cuda.is_available():
+        args.device = "cpu"
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print("Inference completed " f"(patient_ids={num_patient_ids}, rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py
new file mode 100644
index 000000000..558870a63
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py
@@ -0,0 +1,321 @@
+"""Unified multimodal embedding + EHRMamba runner for MIMIC-IV + CXR.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_mamba_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_mamba_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import EHRMamba, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_layers{args.num_layers}"
+        f"_state{args.state_size}"
+        f"_conv{args.conv_kernel}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(), "output", "multimodal_embedding_mamba_mimic4_cxr", run_tag
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = EHRMamba(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        num_layers=args.num_layers,
+        state_size=args.state_size,
+        conv_kernel=args.conv_kernel,
+        dropout=args.dropout,
+        unified_embedding=unified,
+    )
+    print(f"EHRMamba unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + EHRMamba on "
+        "MIMIC-IV mortality with CXR."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--num-layers", type=int, default=2)
+    parser.add_argument("--state-size", type=int, default=16)
+    parser.add_argument("--conv-kernel", type=int, default=4)
+    parser.add_argument("--dropout", type=float, default=0.1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py
new file mode 100644
index 000000000..9ed8467b4
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py
@@ -0,0 +1,325 @@
+"""Unified multimodal embedding + MLP runner for MIMIC-IV + CXR.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_mlp_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_mlp_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import MLP, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_hidden{args.hidden_dim}"
+        f"_layers{args.num_layers}"
+        f"_act{args.activation}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(), "output", "multimodal_embedding_mlp_mimic4_cxr", run_tag
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = MLP(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        hidden_dim=args.hidden_dim,
+        n_layers=args.num_layers,
+        activation=args.activation,
+        unified_embedding=unified,
+    )
+    print(f"MLP unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + MLP on "
+        "MIMIC-IV mortality with CXR."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--hidden-dim", type=int, default=128)
+    parser.add_argument("--num-layers", type=int, default=2)
+    parser.add_argument(
+        "--activation",
+        type=str,
+        default="relu",
+        choices=["relu", "tanh", "sigmoid", "leaky_relu", "elu"],
+    )
+    parser.add_argument("--dropout", type=float, default=0.1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py
new file mode 100644
index 000000000..0a768a0ef
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py
@@ -0,0 +1,326 @@
+"""Unified multimodal embedding + RNN runner for MIMIC-IV + CXR.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_rnn_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_rnn_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import RNN, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_hidden{args.hidden_dim}"
+        f"_{args.rnn_type.lower()}"
+        f"_layers{args.num_layers}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(), "output", "multimodal_embedding_rnn_mimic4_cxr", run_tag
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = RNN(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        hidden_dim=args.hidden_dim,
+        unified_embedding=unified,
+        rnn_type=args.rnn_type,
+        num_layers=args.num_layers,
+        dropout=args.dropout,
+    )
+    print(f"RNN unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + RNN on "
+        "MIMIC-IV mortality with CXR."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--hidden-dim", type=int, default=128)
+    parser.add_argument(
+        "--rnn-type",
+        type=str,
+        default="GRU",
+        choices=["GRU", "LSTM", "RNN"],
+    )
+    parser.add_argument("--num-layers", type=int, default=1)
+    parser.add_argument("--dropout", type=float, default=0.1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py
new file mode 100644
index 000000000..63b7faca9
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py
@@ -0,0 +1,318 @@
+"""Unified multimodal embedding + Transformer runner for MIMIC-IV + CXR.
+
+This script runs chest X-ray (CXR)-only mortality prediction (metadata/negbio)
+with the CXRMIMIC4 task — the imaging-only ablation baseline.
+
+Default roots are set to shared PhysioNet mounts:
+- ehr_root:  /shared/rsaas/physionet.org/files/mimiciv/2.2
+- cxr_root:  /shared/rsaas/physionet.org/files/MIMIC-CXR
+
+Quick start:
+    python examples/mortality_prediction/
+    multimodal_embedding_transformer_mimic4_cxr.py \
+      --quick-test
+
+Smoke test (single forward + inference, no train):
+    python examples/mortality_prediction/
+    multimodal_embedding_transformer_mimic4_cxr.py \
+      --smoke-forward
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import Transformer, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_run_output_path(args: argparse.Namespace) -> str:
+    device_tag = args.device.replace(":", "")
+    lr_tag = f"{args.lr:g}".replace(".", "p")
+    dropout_tag = f"{args.dropout:g}".replace(".", "p")
+    run_tag = (
+        f"cxr{args.cxr_variant}"
+        f"_emb{args.embedding_dim}"
+        f"_layers{args.num_layers}"
+        f"_heads{args.heads}"
+        f"_drop{dropout_tag}"
+        f"_win{args.observation_window_hours}"
+        f"_ep{args.epochs}"
+        f"_bs{args.batch_size}"
+        f"_lr{lr_tag}"
+        f"_{device_tag}"
+        f"_nw{args.num_workers}"
+        f"_seed{args.seed}"
+        f"_dev{int(args.dev)}"
+        f"_quick{int(args.quick_test)}"
+        f"_smoke{int(args.smoke_forward)}"
+    )
+    return os.path.join(
+        os.getcwd(), "output", "multimodal_embedding_transformer_mimic4_cxr", run_tag
+    )
+
+
+def run(args: argparse.Namespace) -> Tuple[int, int]:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+    total_start = time.perf_counter()
+
+    cuda_device_index = None
+    if args.device.startswith("cuda"):
+        device_index = torch.device(args.device).index
+        cuda_device_index = 0 if device_index is None else device_index
+
+    print("Using dataset roots:")
+    print(f"  ehr_root:     {args.ehr_root}")
+    print(f"  cxr_root:     {args.cxr_root}")
+    print(f"  cxr_variant:  {args.cxr_variant}")
+    print(f"  cache_dir:    {args.cache_dir}")
+    print(f"  num_workers:  {args.num_workers}")
+
+    base_dataset = MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        cxr_root=args.cxr_root,
+        cxr_variant=args.cxr_variant,
+        ehr_tables=[],
+        cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+    )
+
+    task = CXRMIMIC4(
+        window_hours=args.observation_window_hours,
+    )
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or use "
+            "--quick-test first."
+        )
+
+    print(f"Task sample count: {len(sample_dataset)}")
+    print("Input processor schemas:")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        if processor is None:
+            print(f"  - {key}: <no processor>")
+            continue
+        print(f"  - {key}: {type(processor).__name__}, " f"schema={processor.schema()}")
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+    model = Transformer(
+        dataset=sample_dataset,
+        embedding_dim=args.embedding_dim,
+        heads=args.heads,
+        dropout=args.dropout,
+        num_layers=args.num_layers,
+        unified_embedding=unified,
+    )
+    print(f"Transformer unified mode: {model._use_unified}")
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    print(
+        "Split sizes: " f"train={len(train_ds)}, val={len(val_ds)}, test={len(test_ds)}"
+    )
+
+    debug_batch = next(iter(train_loader))
+    print("Batch field diagnostics (train batch 0):")
+    for key in sample_dataset.input_schema.keys():
+        processor = sample_dataset.input_processors.get(key)
+        feature = debug_batch.get(key)
+        schema = processor.schema() if processor is not None else ()
+        print(f"  - {key}: type={type(feature).__name__}, schema={schema}")
+
+        if isinstance(feature, tuple):
+            for i, elem in enumerate(feature):
+                shape = getattr(elem, "shape", None)
+                print(f"      tuple[{i}] type={type(elem).__name__} " f"shape={shape}")
+
+        if processor is not None and isinstance(feature, tuple):
+            for field_name in ("value", "time", "mask"):
+                if field_name in schema:
+                    idx = schema.index(field_name)
+                    if idx < len(feature):
+                        selected = feature[idx]
+                        shape = getattr(selected, "shape", None)
+                        print(
+                            f"      schema['{field_name}'] -> tuple[{idx}] "
+                            f"type={type(selected).__name__} shape={shape}"
+                        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc"],
+        device=args.device,
+        enable_logging=True,
+        output_path=_build_run_output_path(args),
+    )
+
+    if not args.smoke_forward and args.epochs > 0 and len(train_ds) > 0:
+        if cuda_device_index is not None:
+            torch.cuda.reset_peak_memory_stats(cuda_device_index)
+            torch.cuda.synchronize(cuda_device_index)
+
+        train_start = time.perf_counter()
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params={"lr": args.lr},
+            monitor=None,
+            load_best_model_at_last=False,
+            max_grad_norm=args.max_grad_norm,
+            accumulation_steps=args.grad_accum,
+            use_amp=args.amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if cuda_device_index is not None:
+            torch.cuda.synchronize(cuda_device_index)
+            peak_train_bytes = torch.cuda.max_memory_allocated(cuda_device_index)
+            peak_train_vram_mb = peak_train_bytes / (1024**2)
+        else:
+            peak_train_vram_mb = None
+        train_runtime_sec = time.perf_counter() - train_start
+    else:
+        peak_train_vram_mb = None
+        train_runtime_sec = None
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+    scores = trainer.evaluate(inference_loader)
+
+    if cuda_device_index is not None:
+        torch.cuda.synchronize(cuda_device_index)
+
+    total_runtime_sec = time.perf_counter() - total_start
+    print("Benchmark summary:")
+    print(f"  total_runtime_sec: {total_runtime_sec:.2f}")
+    if train_runtime_sec is None:
+        print("  training_runtime_sec: N/A (training skipped)")
+        print("  peak_train_vram_mb: N/A (training skipped)")
+    else:
+        print(f"  training_runtime_sec: {train_runtime_sec:.2f}")
+        if peak_train_vram_mb is None:
+            print("  peak_train_vram_mb: N/A (non-CUDA device)")
+        else:
+            print(f"  peak_train_vram_mb: {peak_train_vram_mb:.2f}")
+    print("Evaluation metrics:")
+    for metric_name, metric_value in scores.items():
+        print(f"  {metric_name}: {metric_value:.4f}")
+
+    return len(patient_ids), y_true.shape[0]
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run unified multimodal embedding + Transformer on "
+        "MIMIC-IV mortality with CXR."
+    )
+    parser.add_argument(
+        "--ehr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    )
+    parser.add_argument(
+        "--cxr-root",
+        type=str,
+        default="/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    )
+    parser.add_argument(
+        "--cxr-variant",
+        type=str,
+        default="sunlab",
+        choices=["default", "sunlab"],
+    )
+    parser.add_argument(
+        "--cache-dir",
+        type=str,
+        default="/shared/eng/pyhealth",
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--num-layers", type=int, default=1)
+    parser.add_argument("--heads", type=int, default=1)
+    parser.add_argument("--dropout", type=float, default=0.1)
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=1)
+    parser.add_argument("--lr", type=float, default=1e-3)
+    parser.add_argument("--device", type=str, default="cuda:1")
+    parser.add_argument("--num-workers", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=42)
+
+    parser.add_argument("--dev", action="store_true")
+    parser.add_argument("--quick-test", action="store_true")
+    parser.add_argument("--smoke-forward", action="store_true")
+    parser.add_argument(
+        "--condor",
+        action="store_true",
+        help="Set device to 'cuda' (no index) for HTCondor GPU jobs.",
+    )
+    parser.add_argument("--grad-accum", type=int, default=1)
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+    parser.add_argument("--amp", action="store_true")
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+    )
+    args = parser.parse_args()
+
+    if args.condor:
+        args.device = "cuda"
+    if args.quick_test:
+        args.dev = True
+        args.epochs = 1
+        args.batch_size = min(args.batch_size, 4)
+
+    return args
+
+
+if __name__ == "__main__":
+    cli_args = parse_args()
+    num_patient_ids, num_rows = run(cli_args)
+    print(f"Inference completed (patient_ids={num_patient_ids}, " f"rows={num_rows}).")
diff --git a/examples/mortality_prediction/multimodal_mimic4.py b/examples/mortality_prediction/multimodal_mimic4.py
new file mode 100644
index 000000000..d28bde608
--- /dev/null
+++ b/examples/mortality_prediction/multimodal_mimic4.py
@@ -0,0 +1,134 @@
+import os
+
+# PyHealth Packages
+from pyhealth.datasets import MIMIC4Dataset
+from pyhealth.tasks.multimodal_mimic4 import (
+    ICDLabsMIMIC4,
+    NotesLabsMIMIC4,
+    LabsMIMIC4,
+    CXRMIMIC4,
+)
+
+# Load MIMIC4 Files
+# There's probably better ways dealing with this on the cluster, but working locally for now
+# (see: https://github.com/sunlabuiuc/PyHealth/blob/master/examples/mortality_prediction/multimodal_mimic4_minimal.py)
+
+TASK = "NotesLabsMIMIC4"  # Options: ICDLabsMIMIC4, NotesLabsMIMIC4, LabsMIMIC4, CXRMIMIC4  # Each task isolates a different modality subset so we can evaluate the value of adding more modalities
+DEV_MODE = True
+ENVIRONMENT = "CampusCluster"  # Either 'Local' or 'CampusCluster' or "SunLabCluster"
+NETID = "wp14" # For personal cache
+
+if ENVIRONMENT == "Local":
+    pyhealth_repo_root = "/Users/wpang/Desktop/PyHealth"
+
+    ehr_root = os.path.join(
+        pyhealth_repo_root, "local_data/local/data/physionet.org/files/mimiciv/2.2"
+    )
+    note_root = os.path.join(
+        pyhealth_repo_root,
+        "local_data/local/data/physionet.org/files/mimic-iv-note/2.2",
+    )
+    cxr_root = os.path.join(
+        pyhealth_repo_root,
+        "llocal_data/local/data/physionet.org/files/mimic-cxr-jpg/2.0.0",
+    )
+    cache_dir = os.path.join(
+        pyhealth_repo_root, "local_data/local/data/wp/pyhealth_cache"
+    )
+elif ENVIRONMENT == "CampusCluster":
+
+    ehr_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2"
+    note_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note"
+    cxr_root = None # Please fill this in
+    cache_dir = f"/u/{NETID}/pyhealth_cache"
+elif ENVIRONMENT == "SunLabCluster":
+
+    ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+    note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+    cxr_root = None # Please fill this in
+    cache_dir = f"/home/{NETID}/pyhealth_cache"
+
+
+if __name__ == "__main__":
+
+    if TASK == "ICDLabsMIMIC4":
+        dataset = MIMIC4Dataset(
+            ehr_root=ehr_root,
+            ehr_tables=[
+                "diagnoses_icd",
+                "procedures_icd",
+                "labevents",
+                "prescriptions",
+            ],
+            cache_dir=cache_dir,
+            num_workers=8,
+            dev=DEV_MODE,
+        )
+
+        # Apply multimodal task
+        task = ICDLabsMIMIC4()
+        samples = dataset.set_task(task)
+
+        # Get and print sample
+        sample = samples[0]
+        print(sample)
+
+    elif TASK == "NotesLabsMIMIC4":
+        dataset = MIMIC4Dataset(
+            ehr_root=ehr_root,
+            note_root=note_root,
+            ehr_tables=[
+                "diagnoses_icd",
+                "procedures_icd",
+                "prescriptions",
+                "labevents",
+            ],
+            note_tables=["discharge", "radiology"],
+            cache_dir=cache_dir,
+            num_workers=8,
+            dev=DEV_MODE,
+        )
+
+        # Apply multimodal task
+        task = NotesLabsMIMIC4()
+        samples = dataset.set_task(task)
+
+        # Get and print sample
+        sample = samples[0]
+        print(sample)
+
+    elif TASK == "LabsMIMIC4":
+        dataset = MIMIC4Dataset(
+            ehr_root=ehr_root,
+            ehr_tables=["labevents"],
+            cache_dir=cache_dir,
+            num_workers=8,
+            dev=DEV_MODE,
+        )
+
+        # Apply multimodal task
+        task = LabsMIMIC4()
+        samples = dataset.set_task(task)
+
+        # Get and print sample
+        sample = samples[0]
+        print(sample)
+
+    elif TASK == "CXRMIMIC4":
+        dataset = MIMIC4Dataset(
+            ehr_root=ehr_root,
+            cxr_root=cxr_root,
+            cxr_variant="sunlab",
+            cxr_tables=["metadata", "negbio"],
+            cache_dir=cache_dir,
+            num_workers=8,
+            dev=DEV_MODE,
+        )
+
+        # Apply multimodal task
+        task = CXRMIMIC4()
+        samples = dataset.set_task(task)
+
+        # Get and print sample
+        sample = samples[0]
+        print(sample)
diff --git a/examples/mortality_prediction/readme.md b/examples/mortality_prediction/readme.md
new file mode 100644
index 000000000..fb82f291b
--- /dev/null
+++ b/examples/mortality_prediction/readme.md
@@ -0,0 +1,18 @@
+# Readme file for mortality prediction run examples here.
+
+
+
+
+# Multimodality
+
+## Model Variants
+
+nohup python examples/mortality_prediction/multimodal_embedding_mamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mamba_mimic4_cxr_b1.log &
+
+nohup python examples/mortality_prediction/multimodal_embedding_mlp_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_mlp_mimic4_cxr_b1.log &
+
+nohup python examples/mortality_prediction/multimodal_embedding_rnn_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_rnn_mimic4_cxr_b1.log &
+
+nohup python examples/mortality_prediction/multimodal_embedding_transformer_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_transformer_mimic4_cxr_b1.log &
+
+nohup python examples/mortality_prediction/multimodal_embedding_jamba_mimic4_cxr.py --batch-size 1 --device cuda:3 > ../logs/multimodal_embedding_jamba_mimic4_cxr_b1.log &
\ No newline at end of file
diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
new file mode 100644
index 000000000..58e92ae07
--- /dev/null
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
@@ -0,0 +1,644 @@
+"""End-to-end protocol runner for Unified Embedding on MIMIC-IV.
+
+Trains and evaluates a unified-embedding model (MLP / RNN / Transformer /
+BottleneckTransformer / EHRMamba / JambaEHR) on a MIMIC-IV mortality task,
+then writes per-sample predictions to CSV.
+
+Tasks
+-----
+--task stagenet (default)
+    MortalityPredictionStageNetMIMIC4: ICD codes + 10-dim lab vectors,
+    patient-level samples aggregated across all admissions.
+
+--task icd_labs
+    ICDLabsMIMIC4: ICD codes + 10-dim lab vectors via the unified
+    multimodal pipeline.  No notes required.
+
+--task notes_labs (recommended for multimodal)
+    NotesLabsMIMIC4: admission-context note sections + labs, no ICD codes.
+    Extracts Chief Complaint, HPI, PMH, Medications on Admission from the
+    discharge note — text available at admission time, ~90%+ coverage.
+    Also includes in-window radiology reports (Indication, Impression
+    sections), bounded to the observation window rather than timestamp 0.0.
+    Requires --note-root.
+    Add --freeze-encoder to freeze Bio_ClinicalBERT and train only the
+    backbone; cuts BERT VRAM by ~50%, useful on smaller GPUs (≤24 GB).
+    Add --icd-codes to include discharge-coded ICD codes (ablation only).
+
+--task notes_labs_cxr
+    NotesLabsCXRMIMIC4: same admission-context notes + labs as notes_labs,
+    plus in-window chest X-ray studies. Requires --note-root and --cxr-root.
+
+Example
+-------
+    python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+      --ehr-root /path/to/mimiciv/2.2 \\
+      --task stagenet \\
+      --model transformer \\
+      --heads 4 --num-layers 2 \\
+      --dev --device cpu \\
+      --epochs 10 --batch-size 32 --lr 1e-3 \\
+      --output-dir ./output/unified_e2e
+
+    # EHRMamba on full dataset (no --dev):
+    python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+      --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\
+      --task notes_labs --model ehrmamba \\
+      --embedding-dim 128 --num-layers 2 --seed 42
+
+    # JambaEHR:
+    python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+      --ehr-root /data/mimic-iv/2.2 --note-root /data/mimic-iv/note \\
+      --task notes_labs --model jambaehr \\
+      --embedding-dim 128 --jamba-transformer-layers 2 --jamba-mamba-layers 6
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    sample_balanced,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel
+from pyhealth.models.bottleneck_transformer import BottleneckTransformer
+from pyhealth.models.ehrmamba import EHRMamba
+from pyhealth.models.jamba_ehr import JambaEHR
+from pyhealth.tasks import MortalityPredictionStageNetMIMIC4
+from pyhealth.tasks.multimodal_mimic4 import (
+    ICDLabsMIMIC4,
+    LabsMIMIC4,
+    NotesLabsCXRMIMIC4,
+    NotesLabsMIMIC4,
+)
+from pyhealth.trainer import Trainer
+from pyhealth.utils import set_seed
+
+
+def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
+    ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"]
+    note_tables = None
+    cxr_kwargs = {}
+
+    if args.task == "notes_labs":
+        if not args.note_root:
+            raise ValueError("--task notes_labs requires --note-root.")
+        note_tables = ["discharge", "radiology"]
+        ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] if args.icd_codes else ["labevents"]
+
+    if args.task == "notes_labs_cxr":
+        if not args.note_root:
+            raise ValueError("--task notes_labs_cxr requires --note-root.")
+        if not args.cxr_root:
+            raise ValueError("--task notes_labs_cxr requires --cxr-root.")
+        note_tables = ["discharge", "radiology"]
+        ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"] if args.icd_codes else ["labevents"]
+        cxr_kwargs = dict(
+            cxr_root=args.cxr_root,
+            cxr_variant=args.cxr_variant,
+            cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        )
+
+    if args.task == "icd_labs":
+        ehr_tables = ["diagnoses_icd", "procedures_icd", "labevents"]
+
+    if args.task == "labs":
+        ehr_tables = ["labevents"]
+
+    return MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        ehr_tables=ehr_tables,
+        note_root=args.note_root if note_tables else None,
+        note_tables=note_tables,
+        cache_dir=args.cache_dir,
+        dev=args.dev if args.dev else False,
+        num_workers=args.num_workers,
+        **cxr_kwargs,
+    )
+
+
+def _build_task(args: argparse.Namespace):
+    if args.task == "stagenet":
+        return MortalityPredictionStageNetMIMIC4()
+    if args.task == "icd_labs":
+        return ICDLabsMIMIC4(window_hours=args.observation_window_hours)
+    if args.task == "notes_labs":
+        return NotesLabsMIMIC4(
+            window_hours=args.observation_window_hours,
+            include_icd=args.icd_codes,
+            include_vitals=args.include_vitals,
+        )
+    if args.task == "notes_labs_cxr":
+        return NotesLabsCXRMIMIC4(
+            window_hours=args.observation_window_hours,
+            include_icd=args.icd_codes,
+            include_vitals=args.include_vitals,
+        )
+    if args.task == "labs":
+        return LabsMIMIC4(window_hours=args.observation_window_hours)
+    raise ValueError(f"Unknown task: {args.task}")
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_model(args: argparse.Namespace, sample_dataset: Any):
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+        freeze_text_encoder=args.freeze_encoder,
+    )
+
+    if args.model == "mlp":
+        return MLP(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            hidden_dim=args.hidden_dim,
+            unified_embedding=unified,
+        )
+    if args.model == "rnn":
+        return RNN(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            hidden_dim=args.hidden_dim,
+            unified_embedding=unified,
+            rnn_type=args.rnn_type,
+            num_layers=args.rnn_layers,
+            dropout=args.dropout,
+            bidirectional=args.bidirectional,
+        )
+    if args.model == "transformer":
+        return Transformer(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            heads=args.heads,
+            num_layers=args.num_layers,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "bottleneck_transformer":
+        return BottleneckTransformer(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            bottlenecks_n=args.bottlenecks_n,
+            fusion_startidx=args.fusion_startidx,
+            num_layers=args.num_layers,
+            heads=args.heads,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "ehrmamba":
+        return EHRMamba(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            num_layers=args.num_layers,
+            state_size=args.mamba_state_size,
+            conv_kernel=args.mamba_conv_kernel,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "jambaehr":
+        return JambaEHR(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            num_transformer_layers=args.jamba_transformer_layers,
+            num_mamba_layers=args.jamba_mamba_layers,
+            heads=args.heads,
+            dropout=args.dropout,
+            state_size=args.mamba_state_size,
+            conv_kernel=args.mamba_conv_kernel,
+            unified_embedding=unified,
+        )
+    raise ValueError(f"Unknown model: {args.model}")
+
+
+def _write_predictions(
+    output_csv: Path,
+    patient_ids: list[str],
+    y_true: np.ndarray,
+    y_prob: np.ndarray,
+) -> None:
+    output_csv.parent.mkdir(parents=True, exist_ok=True)
+
+    y_true_flat = y_true.reshape(-1).tolist()
+    y_prob_flat = y_prob.reshape(-1).tolist()
+
+    with output_csv.open("w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(
+            f,
+            fieldnames=["patient_id", "y_true", "y_prob", "y_pred_threshold_0_5"],
+        )
+        writer.writeheader()
+        for idx, prob in enumerate(y_prob_flat):
+            writer.writerow(
+                {
+                    "patient_id": patient_ids[idx],
+                    "y_true": int(y_true_flat[idx]),
+                    "y_prob": float(prob),
+                    "y_pred_threshold_0_5": int(float(prob) >= 0.5),
+                }
+            )
+
+
+def _compute_pos_weight(train_ds, label_key: str = "mortality") -> float:
+    """Count pos/neg in train_ds and return n_neg/n_pos for BCE pos_weight."""
+    n_pos = n_neg = 0
+    for i in range(len(train_ds)):
+        sample = train_ds[i]
+        label = sample.get(label_key, 0)
+        if hasattr(label, "__iter__"):
+            label = next(iter(label))
+        if float(label) > 0.5:
+            n_pos += 1
+        else:
+            n_neg += 1
+    if n_pos == 0:
+        return 1.0
+    # Cap at 10: n_neg/n_pos ≈ 37 on MIMIC-IV mortality is too extreme with
+    # typical LRs and causes training oscillation. 10 still strongly corrects
+    # for imbalance while keeping gradient magnitudes tractable.
+    return min(10.0, n_neg / n_pos)
+
+
+def run(args: argparse.Namespace) -> Path:
+    set_seed(args.seed)
+
+    base_dataset = _build_base_dataset(args)
+    task = _build_task(args)
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or adjust settings."
+        )
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+
+    label_key = list(sample_dataset.output_schema.keys())[0]
+
+    # Resolve effective sampling strategy.
+    # --balanced-sampling / --balanced-ratio are legacy aliases for undersample.
+    strategy = args.sampling_strategy
+    if args.balanced_sampling and strategy == "none":
+        strategy = "undersample"
+
+    if strategy == "undersample":
+        ratio = args.balanced_ratio
+        print(f"[sampling] Undersampling negatives -> pos:neg 1:{ratio}")
+        train_ds = sample_balanced(train_ds, ratio=ratio, seed=args.seed, label_key=label_key)
+        print(f"[sampling] Training size after undersample: {len(train_ds)}")
+
+    model = _build_model(args, sample_dataset)
+
+    # Apply class-imbalance correction via BCE pos_weight.
+    # pos_weight = n_neg / n_pos so the rare positive class gets proportionally
+    # higher gradient signal, preventing all-negative collapse (F1=0).
+    if args.pos_weight is not None:
+        pw_value = args.pos_weight
+    else:
+        print(f"[pos_weight] Computing class balance from {len(train_ds)} training samples...")
+        pw_value = _compute_pos_weight(train_ds, label_key=label_key)
+    print(f"[pos_weight] Using pos_weight={pw_value:.2f} for binary BCE loss.")
+    model._pos_weight = torch.tensor([pw_value], dtype=torch.float32)
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    # Experiment name encodes model + seed for easy log separation
+    exp_name = f"{args.model}_seed{args.seed}"
+    output_dir = Path(args.output_dir)
+
+    wandb_run = None
+    if args.wandb:
+        import wandb
+
+        tags = args.wandb_tags.split(",") if args.wandb_tags else [args.task, args.model]
+        wandb_run = wandb.init(
+            project=args.wandb_project,
+            entity=args.wandb_entity,
+            name=args.wandb_run_name or exp_name,
+            tags=tags,
+            config=vars(args),
+        )
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc", "f1", "accuracy"],
+        device=args.device,
+        enable_logging=True,
+        output_path=str(output_dir),
+        exp_name=exp_name,
+    )
+
+    # BottleneckTransformer is more fragile on full MIMIC-IV with no warmup.
+    # Use safer defaults unless explicitly overridden from CLI.
+    effective_lr = args.lr
+    effective_max_grad_norm = args.max_grad_norm
+    optimizer_params = {}
+
+    if args.model == "bottleneck_transformer":
+        if effective_lr is None:
+            effective_lr = 1e-4
+        if effective_max_grad_norm is None:
+            effective_max_grad_norm = 0.5
+        optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6
+    else:
+        # All non-BT models: 1e-4 (was 1e-3). With pos_weight correction,
+        # effective gradient magnitude for positives is ~10x higher, so a
+        # smaller LR is needed to avoid training oscillation.
+        if effective_lr is None:
+            effective_lr = 1e-4
+        # Universal grad clipping: prevents runaway updates from the weighted
+        # positive-class loss (pos_weight ≈ 10 scales positive gradients 10x).
+        if effective_max_grad_norm is None:
+            effective_max_grad_norm = 1.0
+        if args.adam_eps is not None:
+            optimizer_params["eps"] = args.adam_eps
+
+    optimizer_params["lr"] = effective_lr
+
+    if args.epochs > 0 and len(train_ds) > 0:
+        metrics_history = trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params=optimizer_params,
+            weight_decay=args.weight_decay,
+            max_grad_norm=effective_max_grad_norm,
+            monitor="pr_auc",
+            load_best_model_at_last=True,
+            patience=args.patience,
+            use_amp=args.use_amp,
+            amp_dtype=args.amp_dtype,
+        )
+        if wandb_run is not None:
+            for epoch_record in metrics_history:
+                wandb_run.log(epoch_record, step=epoch_record["epoch"])
+
+    if wandb_run is not None and test_loader is not None:
+        test_scores = trainer.evaluate(test_loader)
+        wandb_run.log({f"test_{k}": v for k, v in test_scores.items()})
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+
+    output_csv = output_dir / exp_name / f"predictions_{args.model}.csv"
+    _write_predictions(output_csv, patient_ids, y_true, y_prob)
+
+    if wandb_run is not None:
+        wandb_run.log({"pos_weight": pw_value})
+        wandb_run.finish()
+
+    return output_csv
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run E2E unified embedding on MIMIC-IV with any of six sequence heads."
+    )
+    parser.add_argument("--ehr-root", type=str, required=True)
+    parser.add_argument("--note-root", type=str, default=None)
+    parser.add_argument("--cxr-root", type=str, default=None)
+    parser.add_argument("--cxr-variant", type=str, default="sunlab", choices=["default", "sunlab"])
+    parser.add_argument("--cache-dir", type=str, default=None)
+    parser.add_argument("--output-dir", type=str, default="./output/unified_e2e")
+
+    parser.add_argument(
+        "--task",
+        type=str,
+        choices=["stagenet", "icd_labs", "labs", "notes_labs", "notes_labs_cxr"],
+        default="stagenet",
+        help=(
+            "notes_labs: admission-context text (CC/HPI/PMH/MedsOnAdm) + labs. "
+            "No ICD codes (discharge-coded = leakage). Recommended for multimodal. "
+            "notes_labs_cxr: notes_labs plus in-window chest X-rays; requires "
+            "--note-root and --cxr-root."
+        ),
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        choices=["mlp", "rnn", "transformer", "bottleneck_transformer",
+                 "ehrmamba", "jambaehr"],
+        default="rnn",
+    )
+
+    # Shared embedding / training
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--hidden-dim", type=int, default=64)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=32)
+    parser.add_argument(
+        "--lr",
+        type=float,
+        default=None,
+        help=(
+            "Learning rate. Default is 1e-4 for all models. "
+            "(Previously 1e-3 for mlp/rnn/transformer/ehrmamba/jambaehr — "
+            "reduced after pos_weight correction caused oscillation at 1e-3.)"
+        ),
+    )
+    parser.add_argument(
+        "--adam-eps",
+        type=float,
+        default=None,
+        help=(
+            "Adam epsilon. Default is model-specific: 1e-8 for non-BT models, "
+            "1e-6 for bottleneck_transformer."
+        ),
+    )
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument(
+        "--use-amp",
+        action="store_true",
+        help="Enable automatic mixed precision training to reduce GPU memory usage.",
+    )
+    parser.add_argument(
+        "--amp-dtype",
+        type=str,
+        default="bf16",
+        choices=["bf16", "fp16"],
+        help="AMP dtype when --use-amp is set. bf16 is more stable (default).",
+    )
+    parser.add_argument("--num-workers", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--patience", type=int, default=None)
+    parser.add_argument(
+        "--dev",
+        nargs="?",
+        type=int,
+        const=1000,
+        default=0,
+        help=(
+            "Dev mode: limit dataset to N patients for fast iteration. "
+            "--dev (no value) defaults to 1000 patients. "
+            "--dev 5000 limits to 5000. Omit for full dataset."
+        ),
+    )
+    parser.add_argument(
+        "--pos-weight",
+        type=float,
+        default=None,
+        help=(
+            "BCE pos_weight for the positive class (float). "
+            "Default: auto-computed as n_neg/n_pos from training split. "
+            "Set to 1.0 to disable class-imbalance correction."
+        ),
+    )
+
+    # Task-specific
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+    parser.add_argument(
+        "--icd-codes",
+        action="store_true",
+        default=False,
+        help=(
+            "Include discharge-coded ICD codes in notes_labs task. "
+            "Default: off (ICD codes are coded at discharge and constitute "
+            "data leakage for in-hospital mortality prediction). "
+            "Enable only for ablation / legacy comparison experiments."
+        ),
+    )
+    parser.add_argument(
+        "--freeze-encoder",
+        action="store_true",
+        default=False,
+        help=(
+            "Freeze pretrained BERT text encoder weights and train only the "
+            "downstream backbone (MLP/RNN/Transformer head + projection layer). "
+            "Reduces VRAM by ~50%% for the text branch; useful when GPU memory "
+            "is limited or for faster iteration on backbone architectures."
+        ),
+    )
+    parser.add_argument(
+        "--include-vitals",
+        action="store_true",
+        default=False,
+        help=(
+            "Include ICU vital signs (HeartRate, SysBP, DiasBP, MeanBP, "
+            "RespRate, SpO2, Temperature) from chartevents as an additional "
+            "modality alongside labs and notes. Adds chartevents to EHR tables."
+        ),
+    )
+    parser.add_argument(
+        "--balanced-sampling",
+        action="store_true",
+        default=False,
+        help=(
+            "Undersample the majority (negative) class in training to improve "
+            "PR-AUC on imbalanced datasets. Uses sample_balanced() to create a "
+            "1:--balanced-ratio pos:neg training set."
+        ),
+    )
+    parser.add_argument(
+        "--balanced-ratio",
+        type=float,
+        default=1.0,
+        help=(
+            "Negatives per positive in the balanced training set. "
+            "Default: 1.0 (equal pos/neg). Only used with --balanced-sampling."
+        ),
+    )
+    parser.add_argument(
+        "--sampling-strategy",
+        type=str,
+        default="none",
+        choices=["none", "undersample"],
+        help=(
+            "Training-set class balance strategy. "
+            "'none': no resampling (default). "
+            "'undersample': drop majority-class (neg) samples via sample_balanced(). "
+            "--balanced-sampling is a legacy alias for 'undersample'."
+        ),
+    )
+
+    # RNN-specific
+    parser.add_argument("--rnn-type", type=str, default="GRU")
+    parser.add_argument("--rnn-layers", type=int, default=1)
+    parser.add_argument("--bidirectional", action="store_true")
+
+    # Transformer / BottleneckTransformer shared
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--num-layers", type=int, default=2)
+
+    # BottleneckTransformer-specific
+    parser.add_argument("--bottlenecks-n", type=int, default=4)
+    parser.add_argument("--fusion-startidx", type=int, default=1)
+
+    # Training stability
+    parser.add_argument(
+        "--max-grad-norm",
+        type=float,
+        default=None,
+        help=(
+            "Gradient clipping max norm. Default is model-specific: None for "
+            "non-BT models, 0.5 for bottleneck_transformer."
+        ),
+    )
+
+    # W&B logging
+    parser.add_argument(
+        "--wandb",
+        action="store_true",
+        default=False,
+        help="Log training/eval metrics to Weights & Biases.",
+    )
+    parser.add_argument("--wandb-project", type=str, default="pyhealth-mortality")
+    parser.add_argument("--wandb-entity", type=str, default=None)
+    parser.add_argument(
+        "--wandb-run-name",
+        type=str,
+        default=None,
+        help="Defaults to '{model}_seed{seed}' if unset.",
+    )
+    parser.add_argument(
+        "--wandb-tags",
+        type=str,
+        default=None,
+        help="Comma-separated wandb tags, e.g. 'labs,rnn'. Defaults to '{task},{model}' if unset.",
+    )
+
+    # Mamba / JambaEHR-specific
+    parser.add_argument("--mamba-state-size", type=int, default=16,
+                        help="SSM state size for EHRMamba and JambaEHR blocks.")
+    parser.add_argument("--mamba-conv-kernel", type=int, default=4,
+                        help="Causal conv kernel size for EHRMamba and JambaEHR blocks.")
+    parser.add_argument("--jamba-transformer-layers", type=int, default=2,
+                        help="Number of Transformer (attention) layers in JambaEHR.")
+    parser.add_argument("--jamba-mamba-layers", type=int, default=6,
+                        help="Number of Mamba (SSM) layers in JambaEHR.")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    output_csv_path = run(args)
+    print(f"Saved predictions to: {output_csv_path}")
diff --git a/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py
new file mode 100644
index 000000000..5e2d0343c
--- /dev/null
+++ b/examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py
@@ -0,0 +1,335 @@
+"""End-to-end protocol runner for Unified Embedding on MIMIC-IV + CXR.
+
+Trains and evaluates a unified-embedding model (MLP / RNN / Transformer /
+BottleneckTransformer / EHRMamba / JambaEHR) on MIMIC-IV mortality using
+chest X-rays only (CXRMIMIC4, the imaging-only ablation baseline), or on
+the ICD + labs stagenet baseline for comparison.
+
+Example
+-------
+    python examples/mortality_prediction/unified_embedding_e2e_mimic4_cxr.py \\
+      --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \\
+      --cxr-root /shared/rsaas/physionet.org/files/MIMIC-CXR \\
+      --task cxr \\
+      --model ehrmamba --num-layers 2 \\
+      --epochs 10 --batch-size 8 --device cuda:0
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+from typing import Any, Optional, Tuple
+
+import numpy as np
+import torch
+
+from pyhealth.datasets import (
+    MIMIC4Dataset,
+    get_dataloader,
+    split_by_patient,
+    split_by_sample,
+)
+from pyhealth.models import MLP, RNN, Transformer, UnifiedMultimodalEmbeddingModel
+from pyhealth.models.bottleneck_transformer import BottleneckTransformer
+from pyhealth.models.ehrmamba import EHRMamba
+from pyhealth.models.jamba_ehr import JambaEHR
+from pyhealth.tasks import MortalityPredictionStageNetMIMIC4
+from pyhealth.tasks.multimodal_mimic4 import CXRMIMIC4
+from pyhealth.trainer import Trainer
+
+
+def _build_base_dataset(args: argparse.Namespace) -> MIMIC4Dataset:
+    ehr_tables = (
+        ["diagnoses_icd", "procedures_icd", "labevents"]
+        if args.task == "stagenet"
+        else []
+    )
+
+    # CXR metadata tables only loaded for the CXR task
+    cxr_kwargs = {}
+    if args.task == "cxr":
+        if not args.cxr_root:
+            raise ValueError("--task cxr requires --cxr-root.")
+        cxr_kwargs = dict(
+            cxr_root=args.cxr_root,
+            cxr_variant=args.cxr_variant,
+            cxr_tables=["metadata", "negbio", "chexpert", "split"],
+        )
+
+    return MIMIC4Dataset(
+        ehr_root=args.ehr_root,
+        ehr_tables=ehr_tables,
+        cache_dir=args.cache_dir,
+        dev=args.dev,
+        num_workers=args.num_workers,
+        **cxr_kwargs,
+    )
+
+
+def _build_task(args: argparse.Namespace):
+    if args.task == "stagenet":
+        return MortalityPredictionStageNetMIMIC4()
+    if args.task == "cxr":
+        return CXRMIMIC4(window_hours=args.observation_window_hours)
+    raise ValueError(f"Unknown task: {args.task}")
+
+
+def _split_dataset(dataset: Any, seed: int) -> Tuple[Any, Any, Any]:
+    train_ds, val_ds, test_ds = split_by_patient(dataset, [0.8, 0.1, 0.1], seed=seed)
+    if len(train_ds) == 0 or len(test_ds) == 0:
+        train_ds, val_ds, test_ds = split_by_sample(dataset, [0.8, 0.1, 0.1], seed=seed)
+    return train_ds, val_ds, test_ds
+
+
+def _build_model(args: argparse.Namespace, sample_dataset: Any):
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=sample_dataset.input_processors,
+        embedding_dim=args.embedding_dim,
+    )
+
+    if args.model == "mlp":
+        return MLP(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            hidden_dim=args.hidden_dim,
+            unified_embedding=unified,
+        )
+    if args.model == "rnn":
+        return RNN(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            hidden_dim=args.hidden_dim,
+            unified_embedding=unified,
+            rnn_type=args.rnn_type,
+            num_layers=args.rnn_layers,
+            dropout=args.dropout,
+            bidirectional=args.bidirectional,
+        )
+    if args.model == "transformer":
+        return Transformer(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            heads=args.heads,
+            num_layers=args.num_layers,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "bottleneck_transformer":
+        return BottleneckTransformer(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            bottlenecks_n=args.bottlenecks_n,
+            fusion_startidx=args.fusion_startidx,
+            num_layers=args.num_layers,
+            heads=args.heads,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "ehrmamba":
+        return EHRMamba(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            num_layers=args.num_layers,
+            state_size=args.mamba_state_size,
+            conv_kernel=args.mamba_conv_kernel,
+            dropout=args.dropout,
+            unified_embedding=unified,
+        )
+    if args.model == "jambaehr":
+        return JambaEHR(
+            dataset=sample_dataset,
+            embedding_dim=args.embedding_dim,
+            num_transformer_layers=args.jamba_transformer_layers,
+            num_mamba_layers=args.jamba_mamba_layers,
+            heads=args.heads,
+            dropout=args.dropout,
+            state_size=args.mamba_state_size,
+            conv_kernel=args.mamba_conv_kernel,
+            unified_embedding=unified,
+        )
+    raise ValueError(f"Unknown model: {args.model}")
+
+
+def _write_predictions(
+    output_csv: Path,
+    patient_ids: list[str],
+    y_true: np.ndarray,
+    y_prob: np.ndarray,
+) -> None:
+    output_csv.parent.mkdir(parents=True, exist_ok=True)
+
+    y_true_flat = y_true.reshape(-1).tolist()
+    y_prob_flat = y_prob.reshape(-1).tolist()
+
+    with output_csv.open("w", newline="", encoding="utf-8") as f:
+        writer = csv.DictWriter(
+            f,
+            fieldnames=["patient_id", "y_true", "y_prob", "y_pred_threshold_0_5"],
+        )
+        writer.writeheader()
+        for idx, prob in enumerate(y_prob_flat):
+            writer.writerow(
+                {
+                    "patient_id": patient_ids[idx],
+                    "y_true": int(y_true_flat[idx]),
+                    "y_prob": float(prob),
+                    "y_pred_threshold_0_5": int(float(prob) >= 0.5),
+                }
+            )
+
+
+def run(args: argparse.Namespace) -> Path:
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    base_dataset = _build_base_dataset(args)
+    task = _build_task(args)
+    sample_dataset = base_dataset.set_task(task, num_workers=args.num_workers)
+
+    if len(sample_dataset) == 0:
+        raise RuntimeError(
+            "Task produced zero samples. Check roots/tables or adjust settings."
+        )
+
+    train_ds, val_ds, test_ds = _split_dataset(sample_dataset, seed=args.seed)
+    model = _build_model(args, sample_dataset)
+
+    train_loader = get_dataloader(train_ds, batch_size=args.batch_size, shuffle=True)
+    val_loader = (
+        get_dataloader(val_ds, batch_size=args.batch_size, shuffle=False)
+        if len(val_ds) > 0
+        else None
+    )
+    test_loader = (
+        get_dataloader(test_ds, batch_size=args.batch_size, shuffle=False)
+        if len(test_ds) > 0
+        else None
+    )
+
+    exp_name = f"{args.model}_seed{args.seed}"
+    output_dir = Path(args.output_dir)
+
+    trainer = Trainer(
+        model=model,
+        metrics=["pr_auc", "roc_auc", "f1", "accuracy"],
+        device=args.device,
+        enable_logging=True,
+        output_path=str(output_dir),
+        exp_name=exp_name,
+    )
+
+    # BT uses tighter optimizer settings to stabilize training on full MIMIC-IV
+    effective_lr = args.lr
+    effective_max_grad_norm = args.max_grad_norm
+    optimizer_params = {}
+
+    if args.model == "bottleneck_transformer":
+        if effective_lr is None:
+            effective_lr = 1e-4
+        if effective_max_grad_norm is None:
+            effective_max_grad_norm = 0.5
+        optimizer_params["eps"] = args.adam_eps if args.adam_eps is not None else 1e-6
+    else:
+        if effective_lr is None:
+            effective_lr = 1e-3
+        if args.adam_eps is not None:
+            optimizer_params["eps"] = args.adam_eps
+
+    optimizer_params["lr"] = effective_lr
+
+    if args.epochs > 0 and len(train_ds) > 0:
+        trainer.train(
+            train_dataloader=train_loader,
+            val_dataloader=val_loader,
+            epochs=args.epochs,
+            optimizer_params=optimizer_params,
+            weight_decay=args.weight_decay,
+            max_grad_norm=effective_max_grad_norm,
+            monitor="pr_auc",
+            load_best_model_at_last=True,
+        )
+
+    inference_loader = test_loader or val_loader or train_loader
+    y_true, y_prob, _, patient_ids = trainer.inference(
+        inference_loader, return_patient_ids=True
+    )
+
+    output_csv = output_dir / exp_name / f"predictions_{args.model}.csv"
+    _write_predictions(output_csv, patient_ids, y_true, y_prob)
+    return output_csv
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Run E2E unified embedding on MIMIC-IV + CXR with any of six backbone models."
+    )
+    parser.add_argument("--ehr-root", type=str, required=True)
+    parser.add_argument("--cxr-root", type=str, default=None)
+    parser.add_argument("--cxr-variant", type=str, default="sunlab", choices=["default", "sunlab"])
+    parser.add_argument("--cache-dir", type=str, default=None)
+    parser.add_argument("--output-dir", type=str, default="./output/unified_e2e_cxr")
+
+    parser.add_argument(
+        "--task", type=str, default="cxr",
+        choices=["stagenet", "cxr"],
+    )
+    parser.add_argument(
+        "--model", type=str, default="rnn",
+        choices=["mlp", "rnn", "transformer", "bottleneck_transformer", "ehrmamba", "jambaehr"],
+    )
+
+    parser.add_argument("--embedding-dim", type=int, default=64)
+    parser.add_argument("--hidden-dim", type=int, default=64)
+    parser.add_argument("--dropout", type=float, default=0.1)
+    parser.add_argument("--epochs", type=int, default=1)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument("--lr", type=float, default=None)
+    parser.add_argument("--adam-eps", type=float, default=None)
+    parser.add_argument("--weight-decay", type=float, default=0.0)
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument("--num-workers", type=int, default=1)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--dev", action="store_true")
+
+    parser.add_argument("--observation-window-hours", type=int, default=24)
+
+    parser.add_argument("--rnn-type", type=str, default="GRU")
+    parser.add_argument("--rnn-layers", type=int, default=1)
+    parser.add_argument("--bidirectional", action="store_true")
+
+    parser.add_argument("--heads", type=int, default=4)
+    parser.add_argument("--num-layers", type=int, default=2)
+
+    parser.add_argument("--bottlenecks-n", type=int, default=4)
+    parser.add_argument("--fusion-startidx", type=int, default=1)
+
+    parser.add_argument("--max-grad-norm", type=float, default=None)
+
+    parser.add_argument("--mamba-state-size", type=int, default=16)
+    parser.add_argument("--mamba-conv-kernel", type=int, default=4)
+    parser.add_argument("--jamba-transformer-layers", type=int, default=2)
+    parser.add_argument("--jamba-mamba-layers", type=int, default=6)
+
+    return parser.parse_args()
+
+
+def print_vram_summary(device: Optional[str] = None) -> None:
+    """Print peak allocated and total VRAM for the given CUDA device."""
+    if not torch.cuda.is_available():
+        print("VRAM summary: CUDA not available.")
+        return
+    idx = torch.device(device).index if device and device.startswith("cuda") else 0
+    if idx is None:
+        idx = 0
+    allocated_mb = torch.cuda.max_memory_allocated(idx) / (1024 ** 2)
+    total_mb = torch.cuda.get_device_properties(idx).total_memory / (1024 ** 2)
+    print(f"VRAM summary (cuda:{idx}): peak_allocated={allocated_mb:.0f} MB / total={total_mb:.0f} MB")
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    output_csv_path = run(args)
+    print(f"Saved predictions to: {output_csv_path}")
+    print_vram_summary(args.device)
diff --git a/examples/multimodal_mimic4_task_tutorial.ipynb b/examples/multimodal_mimic4_task_tutorial.ipynb
new file mode 100644
index 000000000..dece75a07
--- /dev/null
+++ b/examples/multimodal_mimic4_task_tutorial.ipynb
@@ -0,0 +1,1824 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "681e0b4c-eed9-4029-a3ce-9401033a972e",
+   "metadata": {},
+   "source": [
+    "# Multimodal MIMIC-4 Task Tutorial\n",
+    "**Contributors:** William Pang, Joshua Chen, Rian Atri\n",
+    "\n",
+    "This notebook demonstrates how to use the `multimodal_mimic4` task.\n",
+    "\n",
+    "**Related Resources** (*Note: We likely need to update file paths once merged to Sunlab Pyhealth Main*)\n",
+    "- [Multimodal MIMIC4 Example](https://github.com/Multimodal-PyHealth/PyHealth/blob/main/examples/mortality_prediction/multimodal_mimic4.py)\n",
+    "- [Multimodal MIMIC4 Task](https://github.com/Multimodal-PyHealth/PyHealth/blob/main/pyhealth/tasks/multimodal_mimic4.py)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "28226a55-0297-4c2e-b3d3-b06c3e3a9938",
+   "metadata": {},
+   "source": [
+    "## 0. Environment Setup and Loading Packages"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "8477e2ac-322f-49e9-819f-ce7621b214c6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datetime import datetime\n",
+    "from typing import Any, Dict, List, Optional\n",
+    "import os\n",
+    "from pathlib import Path\n",
+    "import tempfile\n",
+    "import shutil \n",
+    "import random\n",
+    "import numpy as np\n",
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "72f279a1-d83c-489b-9a39-9efe89bf434a",
+   "metadata": {},
+   "source": [
+    "*Pyhealth Specific Packages*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "201c3a62-2a2e-47ab-9ab8-eea77812f349",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pyhealth.datasets import MIMIC4Dataset\n",
+    "from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4, ClinicalNotesMIMIC4, ClinicalNotesICDLabsMIMIC4, BaseMultimodalMIMIC4Task"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "12508b1e-b554-41fb-bb21-1ec0a7b435ed",
+   "metadata": {},
+   "source": [
+    "*Set Randomness Seed*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "9a466c4b-2283-443f-94a3-5768babddf50",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "<torch._C.Generator at 0x7feda00afdf0>"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "SEED = 22\n",
+    "random.seed(SEED)\n",
+    "np.random.seed(SEED)\n",
+    "torch.manual_seed(SEED)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "a5f48cc8-28a7-41c8-a90d-d8521f337093",
+   "metadata": {},
+   "source": [
+    "*Mortality Label Filter `(0, 1, or [0, 1])`*"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "518d0bd0-d320-4173-8f92-b4e1abbd9a33",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MORTALITY_LABEL = 0"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e525de6c-0798-4963-b2ed-0e728037a0e9",
+   "metadata": {},
+   "source": [
+    "## 1. Notebook-Specific Utility Functions"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "05ec2dc9-1a0c-4360-a028-d93173dcba65",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def filter_mortality_samples(samples, label=None):\n",
+    "    \"\"\"Return samples filtered by mortality label.\n",
+    "\n",
+    "    Args:\n",
+    "        samples: list of sample dicts with a 'mortality' tensor field.\n",
+    "        label: 0, 1, [0, 1], or None to return all samples.\n",
+    "    \"\"\"\n",
+    "    if label is None or label == [0, 1] or label == [1, 0]:\n",
+    "        return list(samples)\n",
+    "    labels = {label} if isinstance(label, int) else set(label)\n",
+    "    return [s for s in samples if s['mortality'].item() in labels]\n",
+    "\n",
+    "\n",
+    "def print_patient_admission_info(sample):\n",
+    "    patient = dataset.get_patient(sample['patient_id'])\n",
+    "    admissions = patient.get_events(\n",
+    "        event_type=\"admissions\",\n",
+    "        start=sample['window_start'],\n",
+    "        end=sample['window_end'],\n",
+    "    )\n",
+    "    print(f\"\\nnum_admissions:  {len(admissions)}\")\n",
+    "    for adm in admissions:\n",
+    "        print(f\"  hadm_id:       {adm.hadm_id}\")\n",
+    "        print(f\"  admittime:     {adm.timestamp}\")\n",
+    "        print(f\"  dischtime:     {adm.dischtime}\")\n",
+    "\n",
+    "\n",
+    "def print_patient_note_info(sample, note_type=\"discharge\", char_limit=80):\n",
+    "    patient = dataset.get_patient(sample['patient_id'])\n",
+    "    notes = patient.get_events(\n",
+    "        event_type=note_type,\n",
+    "        start=sample['window_start'],\n",
+    "        end=sample['window_end'],\n",
+    "    )\n",
+    "    print(f\"\\n{note_type} notes: {len(notes)}\")\n",
+    "    section_headers = (\n",
+    "        BaseMultimodalMIMIC4Task.DISCHARGE_CLINICAL_HEADERS\n",
+    "        if note_type == \"discharge\"\n",
+    "        else BaseMultimodalMIMIC4Task.RADIOLOGY_CLINICAL_HEADERS\n",
+    "    )\n",
+    "    for note in notes:\n",
+    "        print(f\"  note_id:   {note.note_id}\")\n",
+    "        print(f\"  hadm_id:   {note.hadm_id}\")\n",
+    "        print(f\"  charttime: {note.timestamp}\")\n",
+    "        print(f\"  storetime: {note.storetime}\")\n",
+    "        print(f\"  raw text:      {note.text[:char_limit]}..(Limited to {char_limit} Characters).\")\n",
+    "\n",
+    "        # Show what the task pipeline actually keeps after _parse_note_sections\n",
+    "        parsed = BaseMultimodalMIMIC4Task._parse_note_sections(note.text, note_type=note_type)\n",
+    "        extracted = [f\"{k}: {v}\" for k, v in parsed.items() if k in section_headers and v]\n",
+    "        if extracted:\n",
+    "            filtered_text = \" [SEP] \".join(extracted)\n",
+    "            print(f\"  filtered text: {filtered_text[:char_limit]}..(Limited to {char_limit} Characters).\")\n",
+    "        else:\n",
+    "            print(f\"  filtered text: <no matching section headers {section_headers} - note dropped>\")\n",
+    "\n",
+    "\n",
+    "LABCATEGORIES = {\n",
+    "    itemid: name\n",
+    "    for name, itemids in BaseMultimodalMIMIC4Task.LAB_CATEGORIES.items()\n",
+    "    for itemid in itemids\n",
+    "}\n",
+    "\n",
+    "\n",
+    "def print_patient_lab_info(sample):\n",
+    "    patient = dataset.get_patient(sample['patient_id'])\n",
+    "    labs = patient.get_events(\n",
+    "        event_type=\"labevents\",\n",
+    "        start=sample['window_start'],\n",
+    "        end=sample['window_end'],\n",
+    "    )\n",
+    "    print(f\"\\nlabevents: {len(labs)}\")\n",
+    "    for lab in labs:\n",
+    "        lab_name = LABCATEGORIES.get(lab['itemid'], \"Unknown\")\n",
+    "        print(f\"  itemid:    {lab['itemid']} ({lab_name})\")\n",
+    "        print(f\"  charttime: {lab.timestamp}\")\n",
+    "        print(f\"  storetime: {lab['storetime']}\")\n",
+    "        print(f\"  valuenum:  {lab['valuenum']}\")\n",
+    "\n",
+    "\n",
+    "def print_patient_icd_info(sample):\n",
+    "    patient = dataset.get_patient(sample['patient_id'])\n",
+    "    diagnoses = patient.get_events(\n",
+    "        event_type=\"diagnoses_icd\",\n",
+    "        start=sample['window_start'],\n",
+    "        end=sample['window_end'],\n",
+    "    )\n",
+    "    procedures = patient.get_events(\n",
+    "        event_type=\"procedures_icd\",\n",
+    "        start=sample['window_start'],\n",
+    "        end=sample['window_end'],\n",
+    "    )\n",
+    "    print(f\"\\ndiagnoses_icd: {len(diagnoses)}\")\n",
+    "    for dx in diagnoses:\n",
+    "        print(f\"  hadm_id:  {dx.hadm_id}\")\n",
+    "        print(f\"  seq_num:  {dx.seq_num}\")\n",
+    "        print(f\"  icd_code: {dx.icd_code} (ICD-{dx.icd_version})\\n\")\n",
+    "    print(f\"\\nprocedures_icd: {len(procedures)}\")\n",
+    "    for px in procedures:\n",
+    "        print(f\"  hadm_id:  {px.hadm_id}\")\n",
+    "        print(f\"  seq_num:  {px.seq_num}\")\n",
+    "        print(f\"  icd_code: {px.icd_code} (ICD-{px.icd_version})\\n\")\n",
+    "\n",
+    "def get_project_root(marker=\"PyHealth\"):\n",
+    "    path = Path(os.getcwd())\n",
+    "    for parent in [path, *path.parents]:\n",
+    "        if parent.name == marker:\n",
+    "            return str(parent)\n",
+    "    raise ValueError(f\"'{marker}' not found in path\")\n",
+    "\n",
+    "def clear_cache_directory(path):\n",
+    "    for item in os.listdir(path):\n",
+    "        item_path = os.path.join(path, item)\n",
+    "        if os.path.isdir(item_path):\n",
+    "            shutil.rmtree(item_path)\n",
+    "        else:\n",
+    "            os.remove(item_path)\n",
+    "    print(f\"Cache directory cleared: {path}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "18mf1qapywo",
+   "metadata": {},
+   "source": [
+    "## 2. Load Demo Dataset\n",
+    "\n",
+    "We use the MIMIC-IV demo data in `test-resources/core/mimic4demo`. \n",
+    "\n",
+    "This includes:\n",
+    "- Synthetic Notes (`discharge`, `radiology`)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "tv8ir9mp9t",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "PYHEALTH_REPO_ROOT = get_project_root()\n",
+    "\n",
+    "EHR_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n",
+    "NOTE_ROOT = os.path.join(PYHEALTH_REPO_ROOT, \"test-resources/core/mimic4demo\")\n",
+    "CACHE_DIR = tempfile.mkdtemp()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15ccf548-f1fb-4e51-8c6d-2efcd0e4509c",
+   "metadata": {},
+   "source": [
+    "## 3. Multimodal Variants"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9342654c-da12-466a-86e3-ffc010194c89",
+   "metadata": {},
+   "source": [
+    "### 3.1 ICDLabsMIMIC4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "7ebcce82-abdc-4b17-b3e5-d536fbb97859",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory usage Starting MIMIC4Dataset init: 861.3 MB\n",
+      "Initializing mimic4 dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo|None|None (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193\n",
+      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents', 'prescriptions'] (dev mode: False)\n",
+      "Using default EHR config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
+      "Memory usage Before initializing mimic4_ehr: 861.3 MB\n",
+      "Initializing mimic4_ehr dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/c413fc97-f484-5207-8224-296a85506414\n",
+      "Memory usage After initializing mimic4_ehr: 861.3 MB\n",
+      "Memory usage After EHR dataset initialization: 861.3 MB\n",
+      "Memory usage Completed MIMIC4Dataset init: 861.3 MB\n",
+      "Setting task ICDLabsMIMIC4 for mimic4 base dataset...\n",
+      "Task cache paths: task_df=/tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/tasks/ICDLabsMIMIC4_8f1cf75c-b29c-5bf6-a57e-447d6791544c/task_df.ld, samples=/tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/tasks/ICDLabsMIMIC4_8f1cf75c-b29c-5bf6-a57e-447d6791544c/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n",
+      "Applying task transformations on data with 1 workers...\n",
+      "No cached event dataframe found. Creating: /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/global_event_df.parquet\n",
+      "Combining data from ehr dataset\n",
+      "Scanning table: diagnoses_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: procedures_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: labevents from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n",
+      "Scanning table: prescriptions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/prescriptions.csv.gz\n",
+      "Scanning table: patients from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n",
+      "Scanning table: admissions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: icustays from /home/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n",
+      "Creating combined dataframe\n",
+      "Caching event dataframe to /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/global_event_df.parquet...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 patients. (Polars threads: 128)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['bytes']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████| 13/13 [00:00<00:00, 40.56it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing patients.\n",
+      "Fitting processors on the dataset...\n",
+      "Label mortality vocab: {0: 0, 1: 1}\n",
+      "Processing samples and saving to /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/tasks/ICDLabsMIMIC4_8f1cf75c-b29c-5bf6-a57e-447d6791544c/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n",
+      "Applying processors on data with 1 workers...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 samples. (0 to 13)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['str', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'pickle', 'pickle']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████| 13/13 [00:00<00:00, 2002.42it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing samples.\n",
+      "Cached processed samples to /tmp/tmprh33qezu/c48404f3-04a5-52b8-982a-fb9695894193/tasks/ICDLabsMIMIC4_8f1cf75c-b29c-5bf6-a57e-447d6791544c/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "dataset = MIMIC4Dataset(\n",
+    "    ehr_root=EHR_ROOT,\n",
+    "    ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\", \"prescriptions\"],\n",
+    "    cache_dir=CACHE_DIR,\n",
+    "    num_workers=1,\n",
+    ")\n",
+    "\n",
+    "task = ICDLabsMIMIC4()\n",
+    "samples = dataset.set_task(task, num_workers=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "9b3d9f90-3257-4703-b68c-2ab286ba1e08",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Patient ID:      10002\n",
+      "Mortality Flag:  tensor([0.])\n",
+      "Found 13 unique patient IDs\n",
+      "\n",
+      "num_admissions:  3\n",
+      "  hadm_id:       20000\n",
+      "  admittime:     2151-01-01 07:00:00\n",
+      "  dischtime:     2151-01-02 11:00:00\n",
+      "  hadm_id:       20003\n",
+      "  admittime:     2151-01-10 07:00:00\n",
+      "  dischtime:     2151-01-12 11:00:00\n",
+      "  hadm_id:       20004\n",
+      "  admittime:     2151-04-05 14:00:00\n",
+      "  dischtime:     2151-04-20 09:00:00\n",
+      "\n",
+      "discharge notes: 0\n",
+      "\n",
+      "radiology notes: 0\n",
+      "\n",
+      "labevents: 7\n",
+      "  itemid:    50824 (Sodium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  142.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  4.2\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  24.0\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  145.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  3.9\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  22.0\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  180.0\n",
+      "\n",
+      "diagnoses_icd: 6\n",
+      "  hadm_id:  20000\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1065 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  2\n",
+      "  icd_code: I10 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: J189 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  3\n",
+      "  icd_code: N390 (ICD-10)\n",
+      "\n",
+      "\n",
+      "procedures_icd: 3\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: 3E0234Z (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: 0BH17EZ (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: 5A1935Z (ICD-10)\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n",
+    "random.seed(SEED)\n",
+    "sample = random.choice(samples)\n",
+    "\n",
+    "print(f\"\\nPatient ID:      {sample['patient_id']}\")\n",
+    "print(f\"Mortality Flag:  {sample['mortality']}\")\n",
+    "\n",
+    "# Admissions\n",
+    "print_patient_admission_info(sample)\n",
+    "# Discharge Notes\n",
+    "print_patient_note_info(sample, note_type=\"discharge\")\n",
+    "# Radiology Notes\n",
+    "print_patient_note_info(sample, note_type=\"radiology\")\n",
+    "# Labs\n",
+    "print_patient_lab_info(sample)\n",
+    "# ICD Codes\n",
+    "print_patient_icd_info(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "95c83b49-1903-4042-a3ef-0bf3f1d81b08",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cache directory cleared: /tmp/tmprh33qezu\n"
+     ]
+    }
+   ],
+   "source": [
+    "clear_cache_directory(CACHE_DIR)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "516dbef8-43a6-4ebe-b0d8-29e423b0eb00",
+   "metadata": {},
+   "source": [
+    "### 3.2 ClinicalNotesMIMIC4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "a1c777e0-75a0-4096-9e9b-cb3e955fae46",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory usage Starting MIMIC4Dataset init: 2141.9 MB\n",
+      "Initializing mimic4 dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo|/home/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664\n",
+      "Initializing MIMIC4EHRDataset with tables: [] (dev mode: False)\n",
+      "Using default EHR config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
+      "Memory usage Before initializing mimic4_ehr: 2141.9 MB\n",
+      "Initializing mimic4_ehr dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/dc552f81-8a01-5f67-87b9-cfedce496277\n",
+      "Memory usage After initializing mimic4_ehr: 2141.9 MB\n",
+      "Memory usage After EHR dataset initialization: 2141.9 MB\n",
+      "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n",
+      "Using default note config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n",
+      "Memory usage Before initializing mimic4_note: 2141.9 MB\n",
+      "Initializing mimic4_note dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/ca5f0893-8643-5c09-9196-a135b48792d1\n",
+      "Memory usage After initializing mimic4_note: 2141.9 MB\n",
+      "Memory usage After Note dataset initialization: 2141.9 MB\n",
+      "Memory usage Completed MIMIC4Dataset init: 2141.9 MB\n",
+      "Setting task ClinicalNotesMIMIC4 for mimic4 base dataset...\n",
+      "Task cache paths: task_df=/tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/tasks/ClinicalNotesMIMIC4_aa83d683-2fab-51b6-abbb-c33290f35b7f/task_df.ld, samples=/tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/tasks/ClinicalNotesMIMIC4_aa83d683-2fab-51b6-abbb-c33290f35b7f/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n",
+      "Applying task transformations on data with 1 workers...\n",
+      "No cached event dataframe found. Creating: /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/global_event_df.parquet\n",
+      "Combining data from ehr dataset\n",
+      "Scanning table: patients from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n",
+      "Scanning table: admissions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: icustays from /home/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Combining data from note dataset\n",
+      "Scanning table: discharge from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n",
+      "Scanning table: radiology from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n",
+      "Creating combined dataframe\n",
+      "Caching event dataframe to /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/global_event_df.parquet...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 patients. (Polars threads: 128)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['bytes']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████| 13/13 [00:00<00:00, 91.45it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing patients.\n",
+      "Fitting processors on the dataset...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Label mortality vocab: {0: 0, 1: 1}\n",
+      "Processing samples and saving to /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/tasks/ClinicalNotesMIMIC4_aa83d683-2fab-51b6-abbb-c33290f35b7f/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n",
+      "Applying processors on data with 1 workers...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 samples. (0 to 13)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'no_header_tensor:1', 'pickle', 'pickle']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████| 13/13 [00:00<00:00, 272.26it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing samples.\n",
+      "Cached processed samples to /tmp/tmprh33qezu/fec13f55-d312-5963-a175-d7a9992b1664/tasks/ClinicalNotesMIMIC4_aa83d683-2fab-51b6-abbb-c33290f35b7f/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "dataset = MIMIC4Dataset(\n",
+    "    ehr_root=EHR_ROOT,\n",
+    "    note_root=NOTE_ROOT,\n",
+    "    note_tables=[\"discharge\", \"radiology\"],\n",
+    "    cache_dir=CACHE_DIR,\n",
+    "    num_workers=1,\n",
+    ")\n",
+    "task = ClinicalNotesMIMIC4()\n",
+    "samples = dataset.set_task(task, num_workers=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "570587e6-d3af-444d-85c6-45295a80b6b3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Patient ID:      10002\n",
+      "Mortality Flag:  tensor([0.])\n",
+      "Found 13 unique patient IDs\n",
+      "\n",
+      "num_admissions:  3\n",
+      "  hadm_id:       20000\n",
+      "  admittime:     2151-01-01 07:00:00\n",
+      "  dischtime:     2151-01-02 11:00:00\n",
+      "  hadm_id:       20003\n",
+      "  admittime:     2151-01-10 07:00:00\n",
+      "  dischtime:     2151-01-12 11:00:00\n",
+      "  hadm_id:       20004\n",
+      "  admittime:     2151-04-05 14:00:00\n",
+      "  dischtime:     2151-04-20 09:00:00\n",
+      "\n",
+      "discharge notes: 2\n",
+      "  note_id:   d2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-02 12:00:00\n",
+      "  storetime: 2151-01-02 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: chest pain..(Limited to 80 Characters).\n",
+      "  note_id:   d5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-12 12:00:00\n",
+      "  storetime: 2151-01-12 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: fever and productive cough..(Limited to 80 Characters).\n",
+      "\n",
+      "radiology notes: 3\n",
+      "  note_id:   r2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-01 14:00:00\n",
+      "  storetime: 2151-01-01 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-10 14:00:00\n",
+      "  storetime: 2151-01-10 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r6\n",
+      "  hadm_id:   20004\n",
+      "  charttime: 2151-04-05 14:00:00\n",
+      "  storetime: 2151-04-05 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "\n",
+      "labevents: 0\n",
+      "\n",
+      "diagnoses_icd: 0\n",
+      "\n",
+      "procedures_icd: 0\n"
+     ]
+    }
+   ],
+   "source": [
+    "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n",
+    "random.seed(SEED)\n",
+    "sample = random.choice(samples)\n",
+    "\n",
+    "print(f\"\\nPatient ID:      {sample['patient_id']}\")\n",
+    "print(f\"Mortality Flag:  {sample['mortality']}\")\n",
+    "\n",
+    "# Admissions\n",
+    "print_patient_admission_info(sample)\n",
+    "# Discharge Notes\n",
+    "print_patient_note_info(sample, note_type=\"discharge\")\n",
+    "# Radiology Notes\n",
+    "print_patient_note_info(sample, note_type=\"radiology\")\n",
+    "# Labs\n",
+    "print_patient_lab_info(sample)\n",
+    "# ICD Codes\n",
+    "print_patient_icd_info(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "cd4abcb3-8085-433f-8df6-a190c3a7c491",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cache directory cleared: /tmp/tmprh33qezu\n"
+     ]
+    }
+   ],
+   "source": [
+    "clear_cache_directory(CACHE_DIR)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2697f301-92f3-43a6-9d7d-8beb15fdef24",
+   "metadata": {},
+   "source": [
+    "### 3.3 ClinicalNotesICDLabsMIMIC4"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "546e4811-091b-45c3-97b6-87c79f520f33",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory usage Starting MIMIC4Dataset init: 2240.7 MB\n",
+      "Initializing mimic4 dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo|/home/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e\n",
+      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents', 'prescriptions'] (dev mode: False)\n",
+      "Using default EHR config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
+      "Memory usage Before initializing mimic4_ehr: 2240.7 MB\n",
+      "Initializing mimic4_ehr dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/c413fc97-f484-5207-8224-296a85506414\n",
+      "Memory usage After initializing mimic4_ehr: 2241.0 MB\n",
+      "Memory usage After EHR dataset initialization: 2241.0 MB\n",
+      "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n",
+      "Using default note config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n",
+      "Memory usage Before initializing mimic4_note: 2241.0 MB\n",
+      "Initializing mimic4_note dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/ca5f0893-8643-5c09-9196-a135b48792d1\n",
+      "Memory usage After initializing mimic4_note: 2241.0 MB\n",
+      "Memory usage After Note dataset initialization: 2241.0 MB\n",
+      "Memory usage Completed MIMIC4Dataset init: 2241.0 MB\n",
+      "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n",
+      "Task cache paths: task_df=/tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/task_df.ld, samples=/tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n",
+      "Applying task transformations on data with 1 workers...\n",
+      "No cached event dataframe found. Creating: /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/global_event_df.parquet\n",
+      "Combining data from ehr dataset\n",
+      "Scanning table: diagnoses_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: procedures_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: labevents from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n",
+      "Scanning table: prescriptions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/prescriptions.csv.gz\n",
+      "Scanning table: patients from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n",
+      "Scanning table: admissions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: icustays from /home/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n",
+      "Combining data from note dataset\n",
+      "Scanning table: discharge from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n",
+      "Scanning table: radiology from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n",
+      "Creating combined dataframe\n",
+      "Caching event dataframe to /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/global_event_df.parquet...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 patients. (Polars threads: 128)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['bytes']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████| 13/13 [00:00<00:00, 34.01it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing patients.\n",
+      "Fitting processors on the dataset...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Label mortality vocab: {0: 0, 1: 1}\n",
+      "Processing samples and saving to /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n",
+      "Applying processors on data with 1 workers...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 samples. (0 to 13)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'pickle', 'pickle']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████| 13/13 [00:00<00:00, 307.19it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing samples.\n",
+      "Cached processed samples to /tmp/tmprh33qezu/831d8509-2613-5f7a-91bd-42263d4b180e/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "dataset = MIMIC4Dataset(\n",
+    "    ehr_root=EHR_ROOT,\n",
+    "    note_root=NOTE_ROOT,\n",
+    "    ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\", \"prescriptions\"],\n",
+    "    note_tables=[\"discharge\", \"radiology\"],\n",
+    "    cache_dir=CACHE_DIR,\n",
+    "    num_workers=1,\n",
+    ")\n",
+    "\n",
+    "task = ClinicalNotesICDLabsMIMIC4()\n",
+    "samples = dataset.set_task(task, num_workers=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "id": "544f2f04-22e8-49a2-979d-961ec746fae9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Patient ID:      10002\n",
+      "Mortality Flag:  tensor([0.])\n",
+      "Found 13 unique patient IDs\n",
+      "\n",
+      "num_admissions:  3\n",
+      "  hadm_id:       20000\n",
+      "  admittime:     2151-01-01 07:00:00\n",
+      "  dischtime:     2151-01-02 11:00:00\n",
+      "  hadm_id:       20003\n",
+      "  admittime:     2151-01-10 07:00:00\n",
+      "  dischtime:     2151-01-12 11:00:00\n",
+      "  hadm_id:       20004\n",
+      "  admittime:     2151-04-05 14:00:00\n",
+      "  dischtime:     2151-04-20 09:00:00\n",
+      "\n",
+      "discharge notes: 2\n",
+      "  note_id:   d2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-02 12:00:00\n",
+      "  storetime: 2151-01-02 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: chest pain..(Limited to 80 Characters).\n",
+      "  note_id:   d5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-12 12:00:00\n",
+      "  storetime: 2151-01-12 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: fever and productive cough..(Limited to 80 Characters).\n",
+      "\n",
+      "radiology notes: 3\n",
+      "  note_id:   r2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-01 14:00:00\n",
+      "  storetime: 2151-01-01 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-10 14:00:00\n",
+      "  storetime: 2151-01-10 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r6\n",
+      "  hadm_id:   20004\n",
+      "  charttime: 2151-04-05 14:00:00\n",
+      "  storetime: 2151-04-05 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "\n",
+      "labevents: 7\n",
+      "  itemid:    50824 (Sodium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  142.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  4.2\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  24.0\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  145.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  3.9\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  180.0\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  22.0\n",
+      "\n",
+      "diagnoses_icd: 6\n",
+      "  hadm_id:  20000\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1065 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  2\n",
+      "  icd_code: I10 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  3\n",
+      "  icd_code: N390 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: J189 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "\n",
+      "procedures_icd: 3\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: 3E0234Z (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: 0BH17EZ (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: 5A1935Z (ICD-10)\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n",
+    "random.seed(SEED)\n",
+    "sample = random.choice(samples)\n",
+    "\n",
+    "print(f\"\\nPatient ID:      {sample['patient_id']}\")\n",
+    "print(f\"Mortality Flag:  {sample['mortality']}\")\n",
+    "\n",
+    "# Admissions\n",
+    "print_patient_admission_info(sample)\n",
+    "# Discharge Notes\n",
+    "print_patient_note_info(sample, note_type=\"discharge\")\n",
+    "# Radiology Notes\n",
+    "print_patient_note_info(sample, note_type=\"radiology\")\n",
+    "# Labs\n",
+    "print_patient_lab_info(sample)\n",
+    "# ICD Codes\n",
+    "print_patient_icd_info(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "2c9bec37-ff90-48f2-b340-37232a0cd947",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cache directory cleared: /tmp/tmprh33qezu\n"
+     ]
+    }
+   ],
+   "source": [
+    "clear_cache_directory(CACHE_DIR)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9bce1ec8-2108-4c42-bcc0-35916d05fbed",
+   "metadata": {},
+   "source": [
+    "## 4. Time Filtering"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "3eeef725-c61c-4320-a3bb-0d26c89d693c",
+   "metadata": {},
+   "source": [
+    "### 4.1 Full Patient Window\n",
+    "By default the task uses all data between the first admission time and the last discharge time across the patient's processed admissions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "id": "d3b027ab-c814-49cd-8542-885d76b2401b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory usage Starting MIMIC4Dataset init: 2271.5 MB\n",
+      "Initializing mimic4 dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo|/home/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1\n",
+      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents'] (dev mode: False)\n",
+      "Using default EHR config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
+      "Memory usage Before initializing mimic4_ehr: 2271.5 MB\n",
+      "Initializing mimic4_ehr dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/d9b3e474-d485-53b2-8882-996e19382994\n",
+      "Memory usage After initializing mimic4_ehr: 2271.5 MB\n",
+      "Memory usage After EHR dataset initialization: 2271.5 MB\n",
+      "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n",
+      "Using default note config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n",
+      "Memory usage Before initializing mimic4_note: 2271.5 MB\n",
+      "Initializing mimic4_note dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/ca5f0893-8643-5c09-9196-a135b48792d1\n",
+      "Memory usage After initializing mimic4_note: 2271.5 MB\n",
+      "Memory usage After Note dataset initialization: 2271.5 MB\n",
+      "Memory usage Completed MIMIC4Dataset init: 2271.5 MB\n",
+      "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n",
+      "Task cache paths: task_df=/tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/task_df.ld, samples=/tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n",
+      "Applying task transformations on data with 1 workers...\n",
+      "No cached event dataframe found. Creating: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/global_event_df.parquet\n",
+      "Combining data from ehr dataset\n",
+      "Scanning table: diagnoses_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: procedures_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: labevents from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n",
+      "Scanning table: patients from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n",
+      "Scanning table: admissions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Scanning table: icustays from /home/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n",
+      "Combining data from note dataset\n",
+      "Scanning table: discharge from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n",
+      "Scanning table: radiology from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n",
+      "Creating combined dataframe\n",
+      "Caching event dataframe to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/global_event_df.parquet...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 patients. (Polars threads: 128)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['bytes']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████| 13/13 [00:00<00:00, 37.84it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing patients.\n",
+      "Fitting processors on the dataset...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Label mortality vocab: {0: 0, 1: 1}\n",
+      "Processing samples and saving to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n",
+      "Applying processors on data with 1 workers...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 samples. (0 to 13)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'pickle', 'pickle']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████| 13/13 [00:00<00:00, 302.56it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing samples.\n",
+      "Cached processed samples to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_63249a94-82d6-5f40-a505-b5330ccfcfa7/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "dataset = MIMIC4Dataset(\n",
+    "    ehr_root=EHR_ROOT,\n",
+    "    note_root=NOTE_ROOT,\n",
+    "    ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\"],\n",
+    "    note_tables=[\"discharge\", \"radiology\"],\n",
+    "    cache_dir=CACHE_DIR,\n",
+    "    num_workers=1,\n",
+    ")\n",
+    "\n",
+    "task = ClinicalNotesICDLabsMIMIC4()\n",
+    "samples = dataset.set_task(task, num_workers=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "id": "6b70acb2-bdea-4a43-8b9a-7725c34fcfce",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Patient ID:      10002\n",
+      "Mortality Flag:  tensor([0.])\n",
+      "Found 13 unique patient IDs\n",
+      "\n",
+      "num_admissions:  3\n",
+      "  hadm_id:       20000\n",
+      "  admittime:     2151-01-01 07:00:00\n",
+      "  dischtime:     2151-01-02 11:00:00\n",
+      "  hadm_id:       20003\n",
+      "  admittime:     2151-01-10 07:00:00\n",
+      "  dischtime:     2151-01-12 11:00:00\n",
+      "  hadm_id:       20004\n",
+      "  admittime:     2151-04-05 14:00:00\n",
+      "  dischtime:     2151-04-20 09:00:00\n",
+      "\n",
+      "discharge notes: 2\n",
+      "  note_id:   d2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-02 12:00:00\n",
+      "  storetime: 2151-01-02 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: chest pain..(Limited to 80 Characters).\n",
+      "  note_id:   d5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-12 12:00:00\n",
+      "  storetime: 2151-01-12 13:00:00\n",
+      "  raw text:      Name:  ___                  Unit No:   ___\n",
+      "\n",
+      "Admission Date:  ___              Di..(Limited to 80 Characters).\n",
+      "  filtered text: chief complaint: fever and productive cough..(Limited to 80 Characters).\n",
+      "\n",
+      "radiology notes: 3\n",
+      "  note_id:   r2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-01 14:00:00\n",
+      "  storetime: 2151-01-01 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r5\n",
+      "  hadm_id:   20003\n",
+      "  charttime: 2151-01-10 14:00:00\n",
+      "  storetime: 2151-01-10 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "  note_id:   r6\n",
+      "  hadm_id:   20004\n",
+      "  charttime: 2151-04-05 14:00:00\n",
+      "  storetime: 2151-04-05 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "\n",
+      "labevents: 7\n",
+      "  itemid:    50824 (Sodium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  142.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  4.2\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  24.0\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-01-10 09:00:00\n",
+      "  storetime: 2151-01-10 09:30:00\n",
+      "  valuenum:  145.0\n",
+      "  itemid:    50822 (Potassium)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  3.9\n",
+      "  itemid:    50803 (Bicarbonate)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  22.0\n",
+      "  itemid:    50809 (Glucose)\n",
+      "  charttime: 2151-04-06 08:00:00\n",
+      "  storetime: 2151-04-06 08:30:00\n",
+      "  valuenum:  180.0\n",
+      "\n",
+      "diagnoses_icd: 6\n",
+      "  hadm_id:  20000\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1065 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  2\n",
+      "  icd_code: I10 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  3\n",
+      "  icd_code: N390 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: J189 (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: E1100 (ICD-10)\n",
+      "\n",
+      "\n",
+      "procedures_icd: 3\n",
+      "  hadm_id:  20003\n",
+      "  seq_num:  1\n",
+      "  icd_code: 3E0234Z (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  2\n",
+      "  icd_code: 0BH17EZ (ICD-10)\n",
+      "\n",
+      "  hadm_id:  20004\n",
+      "  seq_num:  1\n",
+      "  icd_code: 5A1935Z (ICD-10)\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n",
+    "random.seed(SEED)\n",
+    "sample = random.choice(samples)\n",
+    "\n",
+    "print(f\"\\nPatient ID:      {sample['patient_id']}\")\n",
+    "print(f\"Mortality Flag:  {sample['mortality']}\")\n",
+    "\n",
+    "# Admissions\n",
+    "print_patient_admission_info(sample)\n",
+    "# Discharge Notes\n",
+    "print_patient_note_info(sample, note_type=\"discharge\")\n",
+    "# Radiology Notes\n",
+    "print_patient_note_info(sample, note_type=\"radiology\")\n",
+    "# Labs\n",
+    "print_patient_lab_info(sample)\n",
+    "# ICD Codes\n",
+    "print_patient_icd_info(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "id": "c6814b1b-f523-4a37-bdef-929440c68952",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cache directory cleared: /tmp/tmprh33qezu\n"
+     ]
+    }
+   ],
+   "source": [
+    "clear_cache_directory(CACHE_DIR)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "fd7d6767-ae51-4278-bb55-d12bdf3a4ff2",
+   "metadata": {},
+   "source": [
+    "### 4.2 Using `window_hours`\n",
+    "By using `window_hours`, we only look at data from `first admission + window hours` to predict the mortality task."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "28b19943-f5ee-4918-9003-60cf3c2ac623",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "WINDOW_HOURS = 24"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "16121916-e52e-4f66-8463-ff0f931cf416",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Memory usage Starting MIMIC4Dataset init: 2301.3 MB\n",
+      "Initializing mimic4 dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo|/home/wp14/PyHealth/test-resources/core/mimic4demo|None (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1\n",
+      "Initializing MIMIC4EHRDataset with tables: ['diagnoses_icd', 'procedures_icd', 'labevents'] (dev mode: False)\n",
+      "Using default EHR config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_ehr.yaml\n",
+      "Memory usage Before initializing mimic4_ehr: 2301.3 MB\n",
+      "Initializing mimic4_ehr dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/d9b3e474-d485-53b2-8882-996e19382994\n",
+      "Memory usage After initializing mimic4_ehr: 2301.3 MB\n",
+      "Memory usage After EHR dataset initialization: 2301.3 MB\n",
+      "Initializing MIMIC4NoteDataset with tables: ['discharge', 'radiology'] (dev mode: False)\n",
+      "Using default note config: /home/wp14/PyHealth/pyhealth/datasets/configs/mimic4_note.yaml\n",
+      "Memory usage Before initializing mimic4_note: 2301.3 MB\n",
+      "Initializing mimic4_note dataset from /home/wp14/PyHealth/test-resources/core/mimic4demo (dev mode: False)\n",
+      "Using provided cache_dir: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/ca5f0893-8643-5c09-9196-a135b48792d1\n",
+      "Memory usage After initializing mimic4_note: 2301.3 MB\n",
+      "Memory usage After Note dataset initialization: 2301.3 MB\n",
+      "Memory usage Completed MIMIC4Dataset init: 2301.3 MB\n",
+      "Setting task ClinicalNotesICDLabsMIMIC4 for mimic4 base dataset...\n",
+      "Task cache paths: task_df=/tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_f3a9a626-ee04-5049-834d-bcda350a0a45/task_df.ld, samples=/tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_f3a9a626-ee04-5049-834d-bcda350a0a45/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n",
+      "Applying task transformations on data with 1 workers...\n",
+      "No cached event dataframe found. Creating: /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/global_event_df.parquet\n",
+      "Combining data from ehr dataset\n",
+      "Scanning table: diagnoses_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/wp14/PyHealth/pyhealth/datasets/mimic4.py:121: UserWarning: Events from discharge table only have date timestamp (no specific time). This may affect temporal ordering of events.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Scanning table: procedures_icd from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/procedures_icd.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: labevents from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/labevents.csv.gz\n",
+      "Joining with table: /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/d_labitems.csv.gz\n",
+      "Scanning table: patients from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/patients.csv.gz\n",
+      "Scanning table: admissions from /home/wp14/PyHealth/test-resources/core/mimic4demo/hosp/admissions.csv.gz\n",
+      "Scanning table: icustays from /home/wp14/PyHealth/test-resources/core/mimic4demo/icu/icustays.csv.gz\n",
+      "Combining data from note dataset\n",
+      "Scanning table: discharge from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/discharge.csv.gz\n",
+      "Scanning table: radiology from /home/wp14/PyHealth/test-resources/core/mimic4demo/note/radiology.csv.gz\n",
+      "Creating combined dataframe\n",
+      "Caching event dataframe to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/global_event_df.parquet...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 patients. (Polars threads: 128)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['bytes']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████| 13/13 [00:00<00:00, 37.50it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing patients.\n",
+      "Fitting processors on the dataset...\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Label mortality vocab: {0: 0, 1: 1}\n",
+      "Processing samples and saving to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_f3a9a626-ee04-5049-834d-bcda350a0a45/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld...\n",
+      "Applying processors on data with 1 workers...\n",
+      "Detected Jupyter notebook environment, setting num_workers to 1\n",
+      "Single worker mode, processing sequentially\n",
+      "Worker 0 started processing 13 samples. (0 to 13)\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "  0%|                                                   | 0/13 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Rank 0 inferred the following `['str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'tensor', 'tensor', 'tensor', 'no_header_tensor:1', 'str', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'tensor', 'no_header_tensor:1', 'pickle', 'pickle']` data format.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████████████████████████████| 13/13 [00:00<00:00, 331.19it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Worker 0 finished processing samples.\n",
+      "Cached processed samples to /tmp/tmprh33qezu/4dc2c58a-f704-5a37-b8bb-cb7d0f4e18a1/tasks/ClinicalNotesICDLabsMIMIC4_f3a9a626-ee04-5049-834d-bcda350a0a45/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "dataset = MIMIC4Dataset(\n",
+    "    ehr_root=EHR_ROOT,\n",
+    "    note_root=NOTE_ROOT,\n",
+    "    ehr_tables=[\"diagnoses_icd\", \"procedures_icd\", \"labevents\"],\n",
+    "    note_tables=[\"discharge\", \"radiology\"],\n",
+    "    cache_dir=CACHE_DIR,\n",
+    "    num_workers=1,\n",
+    ")\n",
+    "\n",
+    "task = ClinicalNotesICDLabsMIMIC4(window_hours=WINDOW_HOURS)\n",
+    "samples = dataset.set_task(task, num_workers=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "42df85ab-2649-497b-92a2-415dc169798e",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Patient ID:      10002\n",
+      "Mortality Flag:  tensor([0.])\n",
+      "Window start:    2151-01-01 07:00:00\n",
+      "Window end:      2151-01-02 07:00:00\n",
+      "Found 13 unique patient IDs\n",
+      "\n",
+      "num_admissions:  1\n",
+      "  hadm_id:       20000\n",
+      "  admittime:     2151-01-01 07:00:00\n",
+      "  dischtime:     2151-01-02 11:00:00\n",
+      "\n",
+      "discharge notes: 0\n",
+      "\n",
+      "radiology notes: 1\n",
+      "  note_id:   r2\n",
+      "  hadm_id:   20000\n",
+      "  charttime: 2151-01-01 14:00:00\n",
+      "  storetime: 2151-01-01 15:00:00\n",
+      "  raw text:      CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagul..(Limited to 80 Characters).\n",
+      "  filtered text: impression: No pneumothorax.  No effusion.  No rib fracture identified...(Limited to 80 Characters).\n",
+      "\n",
+      "labevents: 0\n",
+      "\n",
+      "diagnoses_icd: 0\n",
+      "\n",
+      "procedures_icd: 0\n"
+     ]
+    }
+   ],
+   "source": [
+    "samples = filter_mortality_samples(samples, label=MORTALITY_LABEL)\n",
+    "random.seed(SEED)\n",
+    "sample = random.choice(samples)\n",
+    "\n",
+    "print(f\"\\nPatient ID:      {sample['patient_id']}\")\n",
+    "print(f\"Mortality Flag:  {sample['mortality']}\")\n",
+    "print(f\"Window start:    {sample['window_start']}\")\n",
+    "print(f\"Window end:      {sample['window_end']}\")\n",
+    "\n",
+    "# Admissions\n",
+    "print_patient_admission_info(sample)\n",
+    "# Discharge Notes\n",
+    "print_patient_note_info(sample, note_type=\"discharge\")\n",
+    "# Radiology Notes\n",
+    "print_patient_note_info(sample, note_type=\"radiology\")\n",
+    "# Labs\n",
+    "print_patient_lab_info(sample)\n",
+    "# ICD Codes\n",
+    "print_patient_icd_info(sample)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "57208cc0-94af-456d-9078-12acaf556df7",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Cache directory cleared: /tmp/tmprh33qezu\n"
+     ]
+    }
+   ],
+   "source": [
+    "clear_cache_directory(CACHE_DIR)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/test.py b/examples/test.py
deleted file mode 100644
index 96e785e0b..000000000
--- a/examples/test.py
+++ /dev/null
@@ -1,5 +0,0 @@
-import torch
-
-for i in range(100):
-    a = torch.randn((5, 5, 5)) > 0
-    print(((torch.sum(a, dim=2) != 0) != (torch.any(a != 0, dim=2))).any())
diff --git a/examples/time_image_processor_tutorial.ipynb b/examples/time_image_processor_tutorial.ipynb
index 54d2987dd..3164d0c4b 100644
--- a/examples/time_image_processor_tutorial.ipynb
+++ b/examples/time_image_processor_tutorial.ipynb
@@ -66,7 +66,7 @@
    "source": [
     "## 2. Create Synthetic Time-Stamped X-ray Data\n",
     "\n",
-    "We simulate a scenario where each patient has 1–5 chest X-rays taken at different times during their hospital stay. Each image gets a timestamp representing days from the patient's first admission."
+    "We simulate a scenario where each patient has 1–5 chest X-rays taken at different times during their hospital stay, with one patient having no chest X-rays. Each image gets a timestamp representing days from the patient's first admission."
    ]
   },
   {
@@ -82,8 +82,11 @@
     "NUM_PATIENTS = 20\n",
     "MAX_IMAGES_PER_PATIENT = 5\n",
     "\n",
+    "TOKEN_REPRESENTING_MISSING_PATH = \"\"\n",
+    "TOKEN_REPRESENTING_MISSING_FLOAT = 0.0\n",
+    "\n",
     "samples = []\n",
-    "for pid in range(NUM_PATIENTS):\n",
+    "for pid in range(NUM_PATIENTS-1):\n",
     "    # Each patient has 1-5 X-rays taken at different times\n",
     "    n_images = np.random.randint(1, MAX_IMAGES_PER_PATIENT + 1)\n",
     "\n",
@@ -93,6 +96,11 @@
     "\n",
     "    image_paths = []\n",
     "    for j in range(n_images):\n",
+    "        if pid == np.random.randint(1, NUM_PATIENTS - 1) and j == np.random.randint(0, n_images):\n",
+    "            image_paths.append(TOKEN_REPRESENTING_MISSING_PATH)\n",
+    "            time_diffs[j] = TOKEN_REPRESENTING_MISSING_FLOAT\n",
+    "            continue\n",
+    "        \n",
     "        # Synthetic grayscale X-ray with noise\n",
     "        img_array = np.random.normal(80, 25, (224, 224))\n",
     "\n",
@@ -120,6 +128,13 @@
     "        \"label\": label,\n",
     "    })\n",
     "\n",
+    "# Patient with no CXR at all\n",
+    "samples.append({\n",
+    "    \"patient_id\": f\"p{NUM_PATIENTS-1}\",\n",
+    "    \"visit_id\": f\"v{NUM_PATIENTS-1}\",\n",
+    "    \"chest_xray\": ([TOKEN_REPRESENTING_MISSING_PATH], [TOKEN_REPRESENTING_MISSING_FLOAT]),\n",
+    "    \"label\": label})\n",
+    "\n",
     "print(f\"Created {NUM_PATIENTS} patients in {DATA_ROOT}\")\n",
     "print(f\"Images per patient: 1-{MAX_IMAGES_PER_PATIENT}\")"
    ]
@@ -131,7 +146,9 @@
    "outputs": [],
    "source": [
     "# Inspect a sample patient\n",
-    "sample = samples[0]\n",
+    "PATIENT_NUMBER = 17\n",
+    "\n",
+    "sample = samples[PATIENT_NUMBER]\n",
     "paths, times = sample[\"chest_xray\"]\n",
     "\n",
     "print(f\"Patient {sample['patient_id']}:\")\n",
@@ -166,6 +183,7 @@
     "proc = TimeImageProcessor(\n",
     "    image_size=224,\n",
     "    mode=\"L\",\n",
+    "    padding=TOKEN_REPRESENTING_MISSING_PATH\n",
     ")\n",
     "\n",
     "images, timestamps, tag = proc.process(sample[\"chest_xray\"])\n",
@@ -231,6 +249,7 @@
     "    image_size=224,\n",
     "    mode=\"L\",\n",
     "    max_images=2,\n",
+    "    padding=TOKEN_REPRESENTING_MISSING_PATH\n",
     ")\n",
     "\n",
     "imgs_trunc, ts_trunc, _ = proc_truncated.process(sample[\"chest_xray\"])\n",
@@ -260,6 +279,7 @@
     "proc_norm = TimeImageProcessor(\n",
     "    image_size=128,\n",
     "    mode=\"L\",\n",
+    "    padding=TOKEN_REPRESENTING_MISSING_PATH,\n",
     "    normalize=True,\n",
     "    mean=[0.5],\n",
     "    std=[0.5],\n",
@@ -302,6 +322,7 @@
     "        \"chest_xray\": TimeImageProcessor(\n",
     "            image_size=224,\n",
     "            mode=\"L\",\n",
+    "            padding=TOKEN_REPRESENTING_MISSING_PATH,\n",
     "            max_images=4,\n",
     "        ),\n",
     "    },\n",
@@ -321,7 +342,7 @@
    "outputs": [],
    "source": [
     "# Inspect a processed sample\n",
-    "processed = dataset[0]\n",
+    "processed = dataset[PATIENT_NUMBER]\n",
     "print(f\"Processed sample keys: {list(processed.keys())}\")\n",
     "print()\n",
     "\n",
@@ -353,12 +374,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "proc_demo = TimeImageProcessor(image_size=224, mode=\"L\", max_images=4)\n",
+    "proc_demo = TimeImageProcessor(image_size=224, mode=\"L\", padding=TOKEN_REPRESENTING_MISSING_PATH, max_images=4)\n",
     "\n",
     "print(f\"{'Patient':<10} {'N imgs':<8} {'Output Shape':<25} {'Time Range (days)'}\")\n",
     "print(\"-\" * 65)\n",
     "\n",
-    "for i in range(min(10, len(samples))):\n",
+    "for i in range(len(samples)):\n",
     "    s = samples[i]\n",
     "    paths_i, times_i = s[\"chest_xray\"]\n",
     "    imgs_i, ts_i, _ = proc_demo.process((paths_i, times_i))\n",
@@ -491,13 +512,20 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
   "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
    "name": "python",
-   "version": "3.12.3"
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3"
   }
  },
  "nbformat": 4,
diff --git a/examples/vision_embedding_tutorial.ipynb b/examples/vision_embedding_tutorial.ipynb
index 913422105..d241358c8 100644
--- a/examples/vision_embedding_tutorial.ipynb
+++ b/examples/vision_embedding_tutorial.ipynb
@@ -9,7 +9,7 @@
     "\n",
     "This notebook demonstrates how to use the `VisionEmbeddingModel` for medical imaging tasks in PyHealth.\n",
     "\n",
-    "**Contributors:** Josh Steier \n",
+    "**Contributors:** Josh Steier, Joshua Chen, William Pang\n",
     "\n",
     "\n",
     "**Overview:**\n",
@@ -31,18 +31,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "272cc5bb",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Running on device: cpu\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import os\n",
     "import random\n",
@@ -84,18 +76,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "id": "f07f439c",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Created 200 synthetic images in C:\\Users\\637682\\AppData\\Local\\Temp\\chest_xray_jzp9zbiz\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "USE_SYNTHETIC = True\n",
     "MIMIC_CXR_ROOT = \"/path/to/physionet.org/files/mimic-cxr-jpg/2.0.0\"\n",
@@ -171,22 +155,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "c3f49ff5",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Label label vocab: {0: 0, 1: 1}\n",
-      "Total task samples: 200\n",
-      "Input schema: {'image': 'image'}\n",
-      "Output schema: {'label': 'binary'}\n",
-      "Train/Val/Test sizes: 140, 28, 32\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "image_mode = \"L\" if USE_SYNTHETIC else \"RGB\"\n",
     "\n",
@@ -218,19 +190,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "id": "4684f6c1",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'patient_id': 'list(len=16)', 'visit_id': 'list(len=16)', 'image': 'Tensor(shape=(16, 1, 224, 224))', 'label': 'Tensor(shape=(16, 1))'}\n",
-      "Sample labels: [[0.0], [0.0], [0.0], [1.0], [1.0]]\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "BATCH_SIZE = 16\n",
     "\n",
@@ -260,21 +223,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "id": "a47ebd7e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "patient_id: list(len=16)\n",
-      "visit_id: list(len=16)\n",
-      "image: shape=torch.Size([16, 1, 224, 224]), dtype=torch.float32\n",
-      "label: shape=torch.Size([16, 1]), dtype=torch.float32\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "batch = next(iter(train_loader))\n",
     "\n",
@@ -304,52 +256,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "id": "c4d35bc6",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "C:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth\\sampler\\sage_sampler.py:3: UserWarning: pkg_resources is deprecated as an API. See https://setuptools.pypa.io/en/latest/pkg_resources.html. The pkg_resources package is slated for removal as early as 2025-11-30. Refrain from using this package or pin to Setuptools<81.\n",
-      "  import pkg_resources\n",
-      "c:\\Users\\637682\\Desktop\\pyhealth\\PyHealth\\pyhealth-env\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "VisionEmbeddingModel Output Shape Verification\n",
-      "============================================================\n",
-      "\n",
-      "Input shape: torch.Size([16, 1, 224, 224])  # (B, C, H, W)\n",
-      "\n",
-      "patch:\n",
-      "  Output shape: (16, 197, 128)  # (B, num_tokens, E)\n",
-      "  Tokens: 197 = 196 patches + 1 CLS\n",
-      "  Parameters: 58,240\n",
-      "\n",
-      "cnn:\n",
-      "  Output shape: (16, 50, 128)  # (B, num_tokens, E)\n",
-      "  Tokens: 50 = 49 patches + 1 CLS\n",
-      "  Parameters: 231,808\n",
-      "\n",
-      "resnet18:\n",
-      "  Output shape: (16, 50, 128)  # (B, num_tokens, E)\n",
-      "  Tokens: 50 = 49 patches + 1 CLS\n",
-      "  Parameters: 11,242,432\n",
-      "\n",
-      "resnet50:\n",
-      "  Output shape: (16, 50, 128)  # (B, num_tokens, E)\n",
-      "  Tokens: 50 = 49 patches + 1 CLS\n",
-      "  Parameters: 23,770,560\n",
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from pyhealth.models import VisionEmbeddingModel\n",
     "\n",
@@ -389,6 +299,53 @@
     "    print()"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "604e21d3",
+   "metadata": {},
+   "source": [
+    "## 4.6. Mean Pooling: Compact Single-Vector Embeddings\n",
+    "\n",
+    "Pass `pool='mean'` to collapse all patch tokens into a single averaged vector per image:\n",
+    "\n",
+    "```\n",
+    "Input: (B, C, H, W)\n",
+    "Output with pool='mean': (B, 1, E)\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3be8a484",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"Mean Pooling vs Token Sequence Output\")\n",
+    "print(\"=\" * 60)\n",
+    "print(f\"\\nInput shape: {batch['image'].shape}  # (B, C, H, W)\")\n",
+    "print()\n",
+    "\n",
+    "# With mean pooling: single averaged vector\n",
+    "model_pooled = VisionEmbeddingModel(\n",
+    "    dataset=sample_dataset,\n",
+    "    embedding_dim=128,\n",
+    "    backbone=\"cnn\",\n",
+    "    pool=\"mean\",\n",
+    ")\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    out_pooled = model_pooled({\"image\": batch[\"image\"]})\n",
+    "\n",
+    "info_seq = model_seq.get_output_info(\"image\")\n",
+    "info_pooled = model_pooled.get_output_info(\"image\")\n",
+    "\n",
+    "print(\"With mean pooling (pool='mean'):\")\n",
+    "print(f\"  Output shape: {tuple(out_pooled['image'].shape)}  # (B, 1, E)\")\n",
+    "print(f\"  num_tokens: {info_pooled['num_tokens']}\")\n",
+    "print()"
+   ]
+  },
   {
    "cell_type": "markdown",
    "id": "4ce223e3",
@@ -401,49 +358,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": null,
    "id": "344eac50",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "VisionEmbeddingModel Backbone Comparison\n",
-      "==================================================\n",
-      "\n",
-      "patch:\n",
-      "  Tokens: 197 (196 patches + CLS)\n",
-      "  Parameters: 58,240\n",
-      "\n",
-      "cnn:\n",
-      "  Tokens: 50 (49 patches + CLS)\n",
-      "  Parameters: 231,808\n",
-      "\n",
-      "resnet18:\n",
-      "  Tokens: 50 (49 patches + CLS)\n",
-      "  Parameters: 11,242,432\n",
-      "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to C:\\Users\\637682/.cache\\torch\\hub\\checkpoints\\resnet50-11ad3fa6.pth\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 97.8M/97.8M [00:22<00:00, 4.63MB/s]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\n",
-      "resnet50:\n",
-      "  Tokens: 50 (49 patches + CLS)\n",
-      "  Parameters: 23,770,560\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from pyhealth.models import VisionEmbeddingModel\n",
     "\n",
@@ -476,18 +394,7 @@
    "execution_count": null,
    "id": "91b5f6a6",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Feature keys: ['image']\n",
-      "Label key: label\n",
-      "Mode: binary\n",
-      "Total parameters: 11,259,329\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "class VisionClassifier(nn.Module):\n",
     "    \"\"\"End-to-end classifier using VisionEmbeddingModel.\"\"\"\n",
@@ -575,30 +482,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "id": "699c7e03",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "VisionClassifier(\n",
-      "  (vision_encoder): VisionEmbeddingModel(backbone='resnet18', embedding_dim=128, fields=['image'])\n",
-      "  (classifier): Sequential(\n",
-      "    (0): LayerNorm((128,), eps=1e-05, elementwise_affine=True)\n",
-      "    (1): Linear(in_features=128, out_features=128, bias=True)\n",
-      "    (2): GELU(approximate='none')\n",
-      "    (3): Dropout(p=0.1, inplace=False)\n",
-      "    (4): Linear(in_features=128, out_features=1, bias=True)\n",
-      "  )\n",
-      ")\n",
-      "Metrics: ['roc_auc']\n",
-      "Device: cpu\n",
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from pyhealth.trainer import Trainer\n",
     "\n",
@@ -628,259 +515,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 26,
+   "execution_count": null,
    "id": "b12a766e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Training:\n",
-      "Batch size: 16\n",
-      "Optimizer: <class 'torch.optim.adam.Adam'>\n",
-      "Optimizer params: {'lr': 0.0001}\n",
-      "Weight decay: 0.0\n",
-      "Max grad norm: 1.0\n",
-      "Val dataloader: <torch.utils.data.dataloader.DataLoader object at 0x00000212BB9D8FB0>\n",
-      "Monitor: roc_auc\n",
-      "Monitor criterion: max\n",
-      "Epochs: 5\n",
-      "Patience: None\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3cadbf78b48241199f56d04123c7adb0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Epoch 0 / 5:   0%|          | 0/9 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Train epoch-0, step-9 ---\n",
-      "loss: 0.7026\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.30it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Eval epoch-0, step-9 ---\n",
-      "roc_auc: 0.5000\n",
-      "loss: 0.6933\n",
-      "New best roc_auc score (0.5000) at epoch-0, step-9\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a7bc5082d5b841a2ba394e73c8704263",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Epoch 1 / 5:   0%|          | 0/9 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Train epoch-1, step-18 ---\n",
-      "loss: 0.7016\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.66it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Eval epoch-1, step-18 ---\n",
-      "roc_auc: 0.5000\n",
-      "loss: 0.6935\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e914ba00cdbe4bb0b3d4faa247f07228",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Epoch 2 / 5:   0%|          | 0/9 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Train epoch-2, step-27 ---\n",
-      "loss: 0.6951\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.32it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Eval epoch-2, step-27 ---\n",
-      "roc_auc: 0.5000\n",
-      "loss: 0.6933\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "03347961b139426d8835a8a4e22f5884",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Epoch 3 / 5:   0%|          | 0/9 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Train epoch-3, step-36 ---\n",
-      "loss: 0.6991\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.30it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Eval epoch-3, step-36 ---\n",
-      "roc_auc: 0.5000\n",
-      "loss: 0.6934\n",
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b361d1aff8a744df9fa5a8f12335e81e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Epoch 4 / 5:   0%|          | 0/9 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Train epoch-4, step-45 ---\n",
-      "loss: 0.7007\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  3.76it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "--- Eval epoch-4, step-45 ---\n",
-      "roc_auc: 0.5000\n",
-      "loss: 0.6932\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "trainer.train(\n",
     "    train_dataloader=train_loader,\n",
@@ -901,46 +539,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": null,
    "id": "7dd1577d",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.39it/s]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Validation metrics: roc_auc=0.5000, loss=0.6932\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.41it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Test metrics: roc_auc=0.5000, loss=0.6932\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "evaluation_results = {}\n",
     "for split_name, loader in {\"validation\": val_loader, \"test\": test_loader}.items():\n",
@@ -964,33 +566,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 28,
+   "execution_count": null,
    "id": "5624ba3f",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Evaluation: 100%|██████████| 2/2 [00:00<00:00,  4.39it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Mean loss: 0.6932\n",
-      "Preview (label, positive_prob): [(0.0, 0.49317309260368347), (1.0, 0.49317309260368347), (0.0, 0.49317309260368347), (1.0, 0.49317309260368347), (0.0, 0.49317309260368347)]\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "target_loader = val_loader if val_loader is not None else train_loader\n",
     "\n",
@@ -1011,18 +590,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 29,
+   "execution_count": null,
    "id": "741286d0",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Cleaned up: C:\\Users\\637682\\AppData\\Local\\Temp\\chest_xray_x7mdy6r9\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "if USE_SYNTHETIC:\n",
     "    shutil.rmtree(DATA_ROOT)\n",
@@ -1051,6 +622,7 @@
     "**Key Features:**\n",
     "- Pretrained ImageNet weights (ResNet backbones)\n",
     "- Optional [CLS] token for classification\n",
+    "- `pool='mean'`: collapses patch sequence to a single `(B, 1, E)` vector\n",
     "- Positional embeddings\n",
     "- Compatible with PyHealth's `ImageProcessor`\n",
     "\n",
@@ -1067,9 +639,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "PyHealth Dev",
+   "display_name": "pyhealth",
    "language": "python",
-   "name": "pyhealth-env"
+   "name": "pyhealth"
   },
   "language_info": {
    "codemirror_mode": {
@@ -1081,7 +653,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.3"
+   "version": "3.13.3"
   }
  },
  "nbformat": 4,
diff --git a/pyhealth/data/data.py b/pyhealth/data/data.py
index 14b1b526c..39dd1206f 100644
--- a/pyhealth/data/data.py
+++ b/pyhealth/data/data.py
@@ -155,7 +155,7 @@ def _filter_by_time_range_fast(self, df: pl.DataFrame, start: Optional[datetime]
             start_idx = np.searchsorted(ts_col, np.datetime64(start, "ms"), side="left")
         if end is not None:
             end_idx = np.searchsorted(ts_col, np.datetime64(end, "ms"), side="right")
-        return df.slice(start_idx, end_idx - start_idx)
+        return df.slice(start_idx, max(0, end_idx - start_idx))
 
     def _filter_by_event_type_regular(self, df: pl.DataFrame, event_type: Optional[str]) -> pl.DataFrame:
         """Regular filtering by event type. Time complexity: O(n)."""
diff --git a/pyhealth/datasets/__init__.py b/pyhealth/datasets/__init__.py
index 54e77670c..f2be0499e 100644
--- a/pyhealth/datasets/__init__.py
+++ b/pyhealth/datasets/__init__.py
@@ -58,7 +58,13 @@ def __init__(self, *args, **kwargs):
 from .isruc import ISRUCDataset
 from .medical_transcriptions import MedicalTranscriptionsDataset
 from .mimic3 import MIMIC3Dataset
-from .mimic4 import MIMIC4CXRDataset, MIMIC4Dataset, MIMIC4EHRDataset, MIMIC4NoteDataset
+from .mimic4 import (
+    MIMIC4CXRDataset,
+    MIMIC4CXRSunlabDataset,
+    MIMIC4Dataset,
+    MIMIC4EHRDataset,
+    MIMIC4NoteDataset,
+)
 from .mimicextract import MIMICExtractDataset
 from .omop import OMOPDataset
 from .sample_dataset import SampleBuilder, SampleDataset, create_sample_dataset
diff --git a/pyhealth/datasets/base_dataset.py b/pyhealth/datasets/base_dataset.py
index 0e4280aab..0acae3f39 100644
--- a/pyhealth/datasets/base_dataset.py
+++ b/pyhealth/datasets/base_dataset.py
@@ -3,7 +3,7 @@
 import pickle
 from abc import ABC
 from pathlib import Path
-from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable
+from typing import Dict, Iterator, Iterable, List, Optional, Any, Callable, Union
 import functools
 import operator
 from urllib.parse import urlparse, urlunparse
@@ -84,7 +84,13 @@ def path_exists(path: str) -> bool:
         except requests.RequestException:
             return False
     else:
-        return Path(path).exists()
+        try:
+            return Path(path).exists()
+        except OSError:
+            # Treat unreadable paths (e.g. stale/corrupted filesystem
+            # entries that raise I/O errors on stat) as non-existent so
+            # callers can fall back to an alternate extension.
+            return False
 
 
 def _csv_tsv_gz_path(path: str) -> str:
@@ -314,7 +320,7 @@ class BaseDataset(ABC):
         dataset_name (str): Name of the dataset.
         config (dict): Configuration loaded from a YAML file.
         global_event_df (pl.LazyFrame): The global event data frame.
-        dev (bool): Whether to enable dev mode (limit to 1000 patients).
+        dev (Union[bool, int]): Whether to enable dev mode. If True, limit to 1000 patients. If an int, limit to that many patients.
     """
 
     def __init__(
@@ -325,7 +331,7 @@ def __init__(
         config_path: Optional[str] = None,
         cache_dir: str | Path | None = None,
         num_workers: int = 1,
-        dev: bool = False,
+        dev: Union[bool, int] = False,
     ):
         """Initializes the BaseDataset.
 
@@ -342,7 +348,7 @@ def __init__(
                 - **str** or **Path**: Used as the root cache directory path. A UUID
                   is appended to the provided path to capture dataset configuration.
             num_workers (int): Number of worker processes for parallel operations.
-            dev (bool): Whether to run in dev mode (limits to 1000 patients).
+            dev (Union[bool, int]): Whether to run in dev mode. If True, limits to 1000 patients. If an int, limits to that many patients.
         """
         if len(set(tables)) != len(tables):
             logger.warning("Duplicate table names in tables list. Removing duplicates.")
@@ -500,30 +506,52 @@ def _event_transform(self, output_dir: Path) -> None:
         compute_ok = False
         try:
             df = self.load_data()
-            with DaskCluster(
-                n_workers=self.num_workers,
-                threads_per_worker=1,
-                processes=not in_notebook(),
-                # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory
-                local_directory=str(self.create_tmpdir()),
-            ) as cluster:
-                with DaskClient(cluster) as client:
-                    if self.dev:
-                        logger.info("Dev mode enabled: limiting to 1000 patients")
-                        patients = df["patient_id"].unique().head(1000).tolist()
-                        filter = df["patient_id"].isin(patients)
-                        df = df[filter]
-
-                    logger.info(f"Caching event dataframe to {output_dir}...")
-                    collection = df.sort_values("patient_id").to_parquet(
-                        output_dir,
-                        write_index=False,
-                        compute=False,
-                    )
-                    handle = client.compute(collection)
-                    dask_progress(handle)
-                    handle.result()  # type: ignore
-                    compute_ok = True  # Data is fully written to disk
+            disable_distributed = os.environ.get(
+                "PYHEALTH_DISABLE_DASK_DISTRIBUTED", "0"
+            ) == "1"
+
+            if disable_distributed:
+                logger.info(
+                    "PYHEALTH_DISABLE_DASK_DISTRIBUTED=1 detected; using local dask scheduler."
+                )
+                if self.dev:
+                    n = 1000 if self.dev is True else int(self.dev)
+                    logger.info(f"Dev mode enabled: limiting to {n} patients")
+                    patients = df["patient_id"].unique().head(n, compute=True).tolist()
+                    patient_filter = df["patient_id"].isin(patients)
+                    df = df[patient_filter]
+
+                logger.info(f"Caching event dataframe to {output_dir}...")
+                df.sort_values("patient_id").to_parquet(
+                    output_dir,
+                    write_index=False,
+                    compute=True,
+                )
+            else:
+                with DaskCluster(
+                    n_workers=self.num_workers,
+                    threads_per_worker=1,
+                    processes=not in_notebook(),
+                    # Use cache_dir for Dask's scratch space to avoid filling up /tmp or home directory
+                    local_directory=str(self.create_tmpdir()),
+                ) as cluster:
+                    with DaskClient(cluster) as client:
+                        if self.dev:
+                            logger.info(f"Dev mode enabled: limiting to {1000 if self.dev is True else int(self.dev)} patients")
+                            patients = df["patient_id"].unique().head(1000 if self.dev is True else int(self.dev)).tolist()
+                            filter = df["patient_id"].isin(patients)
+                            df = df[filter]
+
+                        logger.info(f"Caching event dataframe to {output_dir}...")
+                        collection = df.sort_values("patient_id").to_parquet(
+                            output_dir,
+                            write_index=False,
+                            compute=False,
+                        )
+                        handle = client.compute(collection)
+                        dask_progress(handle)
+                        handle.result()  # type: ignore
+                        compute_ok = True  # Data is fully written to disk
         except TimeoutError:
             if compute_ok:
                 # Cluster shutdown timed out after successful compute — data is intact
diff --git a/pyhealth/datasets/collate.py b/pyhealth/datasets/collate.py
index 9e4c113c0..d995dbcdd 100644
--- a/pyhealth/datasets/collate.py
+++ b/pyhealth/datasets/collate.py
@@ -35,6 +35,7 @@ def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]:
 
     * ``Tensor``         — stack if same shape, pad to longest otherwise
     * ``dict[str, ...]`` — recursively collate each sub-key  *(temporal feature)*
+    * ``tuple``          — recursively collate each element assuming homogeneous tuple structure
     * ``int / float``    — ``torch.tensor(...)``
     * anything else      — kept as a plain Python list
 
@@ -69,6 +70,21 @@ def collate_temporal(batch: list[dict[str, Any]]) -> dict[str, Any]:
         elif isinstance(first, torch.Tensor):
             result[key] = _stack_or_pad(vals)
 
+        elif isinstance(first, tuple):
+            # Handle tuple outputs (e.g., from TupleTimeTextProcessor with tokenizer)
+            # Transpose list of tuples to group corresponding elements, then collate each group
+            transposed = list(zip(*vals))  # Now we have [elem0_list, elem1_list, ...]
+            collated_elements = []
+            for elem_vals in transposed:
+                if isinstance(elem_vals[0], torch.Tensor):
+                    if all(e.shape == elem_vals[0].shape for e in elem_vals):
+                        collated_elements.append(torch.stack(list(elem_vals)))
+                    else:
+                        collated_elements.append(pad_sequence(list(elem_vals), batch_first=True))
+                else:
+                    collated_elements.append(list(elem_vals))
+            result[key] = tuple(collated_elements)
+
         elif isinstance(first, (int, float)):
             result[key] = torch.tensor(vals)
 
diff --git a/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml
new file mode 100644
index 000000000..aa4b1c419
--- /dev/null
+++ b/pyhealth/datasets/configs/mimic4_cxr_sunlab.yaml
@@ -0,0 +1,105 @@
+version: "2.1.0"
+tables:
+  metadata:
+    file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+    patient_id: "subject_id"
+    timestamp:
+      - "studydate"
+      - "studytime"
+    timestamp_format: "%Y%m%d%H%M%S"
+    attributes:
+      - "image_path"
+      - "dicom_id"
+      - "study_id"
+      - "performedprocedurestepdescription"
+      - "viewposition"
+      - "rows"
+      - "columns"
+      - "procedurecodesequence_codemeaning"
+      - "viewcodesequence_codemeaning"
+      - "patientorientationcodesequence_codemeaning"
+
+  chexpert:
+    file_path: "mimic-cxr-2.0.0-chexpert.csv"
+    patient_id: "subject_id"
+    join:
+      - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+        "on": "study_id"
+        how: "inner"
+        columns:
+          - "studydate"
+          - "studytime"
+          - "dicom_id"
+    timestamp:
+      - "studydate"
+      - "studytime"
+    timestamp_format: "%Y%m%d%H%M%S"
+    attributes:
+      - "dicom_id"
+      - "study_id"
+      - "atelectasis"
+      - "cardiomegaly"
+      - "consolidation"
+      - "edema"
+      - "enlarged cardiomediastinum"
+      - "fracture"
+      - "lung lesion"
+      - "lung opacity"
+      - "no finding"
+      - "pleural effusion"
+      - "pleural other"
+      - "pneumonia"
+      - "pneumothorax"
+      - "support devices"
+
+  negbio:
+    file_path: "mimic-cxr-2.0.0-negbio.csv"
+    patient_id: "subject_id"
+    join:
+      - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+        "on": "study_id"
+        how: "inner"
+        columns:
+          - "studydate"
+          - "studytime"
+          - "dicom_id"
+    timestamp:
+      - "studydate"
+      - "studytime"
+    timestamp_format: "%Y%m%d%H%M%S"
+    attributes:
+      - "dicom_id"
+      - "study_id"
+      - "atelectasis"
+      - "cardiomegaly"
+      - "consolidation"
+      - "edema"
+      - "enlarged cardiomediastinum"
+      - "fracture"
+      - "lung lesion"
+      - "lung opacity"
+      - "no finding"
+      - "pleural effusion"
+      - "pleural other"
+      - "pneumonia"
+      - "pneumothorax"
+      - "support devices"
+
+  split:
+    file_path: "mimic-cxr-2.0.0-split.csv"
+    patient_id: "subject_id"
+    join:
+      - file_path: "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"
+        "on": "dicom_id"
+        how: "inner"
+        columns:
+          - "studydate"
+          - "studytime"
+    timestamp:
+      - "studydate"
+      - "studytime"
+    timestamp_format: "%Y%m%d%H%M%S"
+    attributes:
+      - "dicom_id"
+      - "study_id"
+      - "split"
diff --git a/pyhealth/datasets/configs/mimic4_ehr.yaml b/pyhealth/datasets/configs/mimic4_ehr.yaml
index 84c570bb9..414f31308 100644
--- a/pyhealth/datasets/configs/mimic4_ehr.yaml
+++ b/pyhealth/datasets/configs/mimic4_ehr.yaml
@@ -108,6 +108,28 @@ tables:
       - "flag"
       - "storetime"
 
+  chartevents:
+    file_path: "icu/chartevents.csv.gz"
+    patient_id: "subject_id"
+    join:
+      - file_path: "icu/d_items.csv.gz"
+        "on": "itemid"
+        how: "inner"
+        columns:
+          - "label"
+          - "category"
+    timestamp: "charttime"
+    attributes:
+      - "hadm_id"
+      - "stay_id"
+      - "itemid"
+      - "label"
+      - "category"
+      - "value"
+      - "valuenum"
+      - "valueuom"
+      - "storetime"
+
   hcpcsevents:
     file_path: "hosp/hcpcsevents.csv.gz"
     patient_id: "subject_id"
diff --git a/pyhealth/datasets/mimic4.py b/pyhealth/datasets/mimic4.py
index 9d1aa55d8..089a6d2f3 100644
--- a/pyhealth/datasets/mimic4.py
+++ b/pyhealth/datasets/mimic4.py
@@ -1,7 +1,7 @@
 import logging
 import os
 import warnings
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import pandas as pd
 import dask.dataframe as dd
@@ -223,6 +223,102 @@ def process_image_path(x):
         return
 
 
+class MIMIC4CXRSunlabDataset(BaseDataset):
+    """
+    Sunlab variant of the MIMIC-CXR Chest X-ray dataset.
+
+    This variant uses the existing metadata CSV and derives flattened image
+    paths at ``images/{dicom_id}.jpg``.
+    """
+
+    def __init__(
+        self,
+        root: str,
+        tables: List[str],
+        dataset_name: str = "mimic4_cxr_sunlab",
+        config_path: Optional[str] = None,
+        cache_dir: Optional[str] = None,
+        **kwargs,
+    ):
+        if config_path is None:
+            config_path = os.path.join(
+                os.path.dirname(__file__), "configs", "mimic4_cxr_sunlab.yaml"
+            )
+            logger.info(f"Using default Sunlab CXR config: {config_path}")
+        self.prepare_metadata(root)
+        log_memory_usage(f"Before initializing {dataset_name}")
+        super().__init__(
+            root=root,
+            tables=tables,
+            dataset_name=dataset_name,
+            config_path=config_path,
+            cache_dir=cache_dir,
+            **kwargs,
+        )
+        log_memory_usage(f"After initializing {dataset_name}")
+
+    @staticmethod
+    def _resolve_column_name(columns: List[str], target: str) -> str:
+        lower_to_original = {col.lower(): col for col in columns}
+        resolved = lower_to_original.get(target.lower())
+        if resolved is None:
+            raise ValueError(
+                f"Expected column '{target}' in metadata, available columns: {columns}"
+            )
+        return resolved
+
+    def prepare_metadata(self, root: str) -> None:
+        metadata_path = os.path.join(root, "mimic-cxr-2.0.0-metadata.csv")
+        if not os.path.exists(metadata_path):
+            raise FileNotFoundError(
+                f"Sunlab metadata file not found: {metadata_path}. "
+                "Expected existing metadata linked by dicom_id/subject_id/study_id."
+            )
+
+        images_dir = os.path.join(root, "images")
+        if not os.path.isdir(images_dir):
+            raise FileNotFoundError(
+                f"Sunlab images directory not found: {images_dir}. "
+                "Expected flattened image files at images/{dicom_id}.jpg."
+            )
+
+        metadata = pd.read_csv(metadata_path, dtype=str)
+
+        dicom_col = self._resolve_column_name(metadata.columns.tolist(), "dicom_id")
+        study_time_col = self._resolve_column_name(
+            metadata.columns.tolist(), "studytime"
+        )
+
+        # Normalize StudyTime so timestamps parse with %Y%m%d%H%M%S in config.
+        def normalize_studytime(value: Optional[str]) -> str:
+            if value is None:
+                return "000000"
+            value_str = str(value).strip()
+            if value_str == "" or value_str.lower() == "nan":
+                return "000000"
+            try:
+                return f"{int(float(value_str)):06d}"
+            except Exception:
+                digits = "".join(ch for ch in value_str if ch.isdigit())
+                if digits == "":
+                    return "000000"
+                return digits[:6].zfill(6)
+
+        metadata[study_time_col] = metadata[study_time_col].apply(normalize_studytime)
+
+        metadata["image_path"] = metadata[dicom_col].apply(
+            lambda dicom_id: os.path.join(root, "images", f"{dicom_id}.jpg")
+        )
+
+        # Align with existing config conventions by using lowercase headers.
+        metadata.columns = [col.lower() for col in metadata.columns]
+
+        metadata.to_csv(
+            os.path.join(root, "mimic-cxr-2.0.0-metadata-pyhealth-sunlab.csv"),
+            index=False,
+        )
+
+
 class MIMIC4Dataset(BaseDataset):
     """
     Unified MIMIC-IV dataset with support for EHR, clinical notes, and X-rays.
@@ -242,6 +338,7 @@ class MIMIC4Dataset(BaseDataset):
         ehr_config_path: Path to the EHR config file
         note_config_path: Path to the note config file
         cxr_config_path: Path to the CXR config file
+        cxr_variant: Which CXR variant to load ("default" or "sunlab")
         dataset_name: Name of the dataset
         dev: Whether to enable dev mode (limit to 1000 patients)
 
@@ -279,8 +376,9 @@ def __init__(
         ehr_config_path: Optional[str] = None,
         note_config_path: Optional[str] = None,
         cxr_config_path: Optional[str] = None,
+        cxr_variant: str = "default",
         dataset_name: str = "mimic4",
-        dev: bool = False,
+        dev: Union[bool, int] = False,
         cache_dir: Optional[str] = None,
         num_workers: int = 1,
     ):
@@ -340,17 +438,33 @@ def __init__(
 
         # Initialize CXR dataset if root is provided
         if cxr_root is not None:
+            if cxr_variant not in {"default", "sunlab"}:
+                raise ValueError(
+                    f"Unknown cxr_variant '{cxr_variant}'. "
+                    "Expected one of {'default', 'sunlab'}."
+                )
+
             logger.info(
-                f"Initializing MIMIC4CXRDataset with tables: {cxr_tables} (dev mode: {dev})"
-            )
-            self.sub_datasets["cxr"] = MIMIC4CXRDataset(
-                root=cxr_root,
-                tables=cxr_tables,
-                config_path=cxr_config_path,
-                cache_dir=str(self.cache_dir),
-                dev=dev,
-                num_workers=num_workers,
+                f"Initializing MIMIC4 CXR variant '{cxr_variant}' with tables: {cxr_tables} (dev mode: {dev})"
             )
+            if cxr_variant == "sunlab":
+                self.sub_datasets["cxr"] = MIMIC4CXRSunlabDataset(
+                    root=cxr_root,
+                    tables=cxr_tables,
+                    config_path=cxr_config_path,
+                    cache_dir=str(self.cache_dir),
+                    dev=dev,
+                    num_workers=num_workers,
+                )
+            else:
+                self.sub_datasets["cxr"] = MIMIC4CXRDataset(
+                    root=cxr_root,
+                    tables=cxr_tables,
+                    config_path=cxr_config_path,
+                    cache_dir=str(self.cache_dir),
+                    dev=dev,
+                    num_workers=num_workers,
+                )
             log_memory_usage("After CXR dataset initialization")
 
         log_memory_usage("Completed MIMIC4Dataset init")
diff --git a/pyhealth/datasets/splitter.py b/pyhealth/datasets/splitter.py
index 2dbc94186..b2ea98854 100644
--- a/pyhealth/datasets/splitter.py
+++ b/pyhealth/datasets/splitter.py
@@ -22,6 +22,7 @@ def sample_balanced(
     ratio: float = 1.0,
     subsample: float = 1.0,
     seed: Optional[int] = None,
+    label_key: str = "label",
 ) -> SampleDataset:
     """Keep positives and negatives at a target ratio, then cap total size.
 
@@ -32,6 +33,7 @@ def sample_balanced(
             exceeds ``len(dataset) * subsample``, both positives and negatives are downsampled
             proportionally while preserving the ratio as closely as possible.
         seed: Optional RNG seed for reproducible negative sampling.
+        label_key: Key to use for accessing the label field in each sample. Default ``"label"``.
 
     Returns:
         A new ``SampleDataset`` containing all positives plus sampled negatives,
@@ -49,7 +51,7 @@ def sample_balanced(
     neg_indices: List[int] = []
 
     for idx in range(len(dataset)):
-        label = _label_to_int(dataset[idx]["label"])
+        label = _label_to_int(dataset[idx][label_key])
         if label == 1:
             pos_indices.append(idx)
         else:
diff --git a/pyhealth/datasets/utils.py b/pyhealth/datasets/utils.py
index 24c87a1d5..ce958f627 100644
--- a/pyhealth/datasets/utils.py
+++ b/pyhealth/datasets/utils.py
@@ -15,9 +15,10 @@
 
 MODULE_CACHE_PATH = os.path.join(BASE_CACHE_PATH, "datasets")
 create_directory(MODULE_CACHE_PATH)
-#PyG import for graph-based models
+# PyG import for graph-based models
 try:
     from torch_geometric.data import Data as PyGData, Batch as PyGBatch
+
     HAS_PYG = True
 except ImportError:
     HAS_PYG = False
@@ -267,40 +268,37 @@ def collate_fn_dict_with_padding(batch: List[dict]) -> dict:
     for key in keys:
         values = [sample[key] for sample in batch]
 
-        # Check if this is a temporal feature tuple (time, values)
-        if isinstance(values[0], tuple) and len(values[0]) == 2:
-            # Handle (time, values) tuples from processors
-            time_tensors = [v[0] for v in values]
-            value_tensors = [v[1] for v in values]
-
-            # Collate values
-            if value_tensors[0].dim() == 0:
-                # Scalars
-                collated_values = torch.stack(value_tensors)
-            elif all(v.shape == value_tensors[0].shape for v in value_tensors):
-                # All same shape
-                collated_values = torch.stack(value_tensors)
-            else:
-                # Variable shapes, use pad_sequence
-                collated_values = pad_sequence(
-                    value_tensors, batch_first=True, padding_value=0
-                )
-
-            # Collate times (if present)
-            collated_times = None
-            # Check if ALL samples have time (not just some)
-            if all(t is not None for t in time_tensors):
-                time_tensors_all = [t for t in time_tensors if t is not None]
-                if all(t.shape == time_tensors_all[0].shape for t in time_tensors_all):
-                    collated_times = torch.stack(time_tensors_all)
+        if isinstance(values[0], tuple):
+            # Generic tuple collation for processor outputs, e.g.
+            # - (time, value) from StageNet processors
+            # - (value, mask, token_type_ids, time, type_tag)
+            #   from TupleTimeTextProcessor with tokenizer.
+            transposed = list(zip(*values))
+            collated_elems = []
+
+            for elem_vals in transposed:
+                first = elem_vals[0]
+
+                if first is None and all(v is None for v in elem_vals):
+                    collated_elems.append(None)
+                elif isinstance(first, torch.Tensor):
+                    tensor_vals = list(elem_vals)
+                    if all(v.shape == tensor_vals[0].shape for v in tensor_vals):
+                        collated_elems.append(torch.stack(tensor_vals))
+                    else:
+                        collated_elems.append(
+                            pad_sequence(
+                                tensor_vals,
+                                batch_first=True,
+                                padding_value=0,
+                            )
+                        )
                 else:
-                    collated_times = pad_sequence(
-                        time_tensors_all, batch_first=True, padding_value=0
-                    )
+                    collated_elems.append(list(elem_vals))
+
+            collated[key] = tuple(collated_elems)
 
-            # Return as tuple (time, values)
-            collated[key] = (collated_times, collated_values)
-            # PyG Data objects (graph processor output)
+        # PyG Data objects (graph processor output)
         elif HAS_PYG and isinstance(values[0], PyGData):
             collated[key] = PyGBatch.from_data_list(values)
 
diff --git a/pyhealth/models/__init__.py b/pyhealth/models/__init__.py
index 5233b1726..3cb4a59af 100644
--- a/pyhealth/models/__init__.py
+++ b/pyhealth/models/__init__.py
@@ -1,12 +1,27 @@
 from .adacare import AdaCare, AdaCareLayer, MultimodalAdaCare
 from .agent import Agent, AgentLayer
 from .base_model import BaseModel
+from .bottleneck_transformer import (
+    BottleneckTransformer,
+    MultimodalBottleneckTransformerEncoder
+)
 from .biot import BIOT
 from .cnn import CNN, CNNLayer
 from .concare import ConCare, ConCareLayer
 from .contrawr import ContraWR, ResBlock2D
 from .deepr import Deepr, DeeprLayer
-from .embedding import EmbeddingModel
+from .embedding import (
+    BaseEmbeddingModel,
+    EmbeddingModel,
+    VisionEmbeddingModel,
+    TextEmbeddingModel,
+    TextEmbedding,  # backward compat alias
+    UnifiedMultimodalEmbeddingModel,
+    SinusoidalTimeEmbedding,
+    PatchEmbedding,
+    Permute,
+    init_embedding_with_pretrained,
+)
 from .gamenet import GAMENet, GAMENetLayer
 from .jamba_ehr import JambaEHR, JambaLayer
 from .logistic_regression import LogisticRegression
@@ -39,8 +54,4 @@
 from .transformers_model import TransformersModel
 from .ehrmamba import EHRMamba, MambaBlock
 from .vae import VAE
-from .vision_embedding import VisionEmbeddingModel
-from .text_embedding import TextEmbedding
 from .sdoh import SdohClassifier
-from .medlink import MedLink
-from .unified_embedding import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding
diff --git a/pyhealth/models/bottleneck_transformer.py b/pyhealth/models/bottleneck_transformer.py
new file mode 100644
index 000000000..61a360ee1
--- /dev/null
+++ b/pyhealth/models/bottleneck_transformer.py
@@ -0,0 +1,520 @@
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
+
+import torch
+import torch.nn as nn
+
+from pyhealth.datasets import SampleDataset
+from pyhealth.models import BaseModel
+from pyhealth.models.embedding import EmbeddingModel
+from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel
+
+
+class MultimodalBottleneckTransformerEncoder(nn.Module):
+    """
+    Generalized Bottleneck Transformer Encoder for N modalities.
+    Based on "Attention Bottlenecks for Multimodal Fusion" (Nagrani et al., NeurIPS 2021).
+    """
+
+    def __init__(
+        self,
+        n_modality: int,
+        bottlenecks_n: int,
+        fusion_startidx: int,
+        n_layers: int,
+        n_head: int,
+        d_model: int,
+        d_ff: int,
+        dropout: float = 0.1,
+    ):
+        super(MultimodalBottleneckTransformerEncoder, self).__init__()
+
+        self.n_modality = n_modality
+        self.fusion_startidx = fusion_startidx
+        self.n_layers = n_layers
+        self.n_fusion_layers = n_layers - fusion_startidx
+        self.n_prefusion = fusion_startidx
+        self.d_model = d_model
+        self.n_bottlenecks = bottlenecks_n
+
+        # Shared Bottleneck Tokens — small init to avoid early gradient explosion
+        self.bottlenecks = nn.Parameter(torch.randn(1, bottlenecks_n, d_model) * 0.02)
+
+        # Prefusion Stacks: independent layers per modality
+        self.prefusion_stacks = nn.ModuleList([
+            nn.ModuleList([
+                nn.TransformerEncoderLayer(
+                    d_model=d_model,
+                    nhead=n_head,
+                    dim_feedforward=d_ff,
+                    dropout=dropout,
+                    batch_first=True
+                ) for _ in range(n_modality)
+            ]) for _ in range(self.n_prefusion)
+        ])
+
+        # Fusion Stacks: processes [bottleneck_tokens || modality_tokens]
+        self.fusion_stacks = nn.ModuleList([
+            nn.ModuleList([
+                nn.TransformerEncoderLayer(
+                    d_model=d_model,
+                    nhead=n_head,
+                    dim_feedforward=d_ff,
+                    dropout=dropout,
+                    batch_first=True
+                ) for _ in range(n_modality)
+            ]) for _ in range(self.n_fusion_layers)
+        ])
+
+    def forward_prefusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]:
+        for enc_layers in self.prefusion_stacks:
+            enc_outputs = []
+            for modal_idx, enc_layer in enumerate(enc_layers):
+                # Apply mask to padding tokens (src_key_padding_mask requires True for ignoring)
+                # True in mask = invalid/padding
+                enc_out = enc_layer(enc_inputs[modal_idx], src_key_padding_mask=~masks[modal_idx] if masks[modal_idx] is not None else None)
+                enc_outputs.append(enc_out)
+            enc_inputs = enc_outputs
+        return enc_inputs
+
+    def forward_fusion(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor], bottleneck_tokens: torch.Tensor, valid_modalities: List[torch.Tensor]) -> List[torch.Tensor]:
+        # valid_modalities: [B] list of boolean/float tensors indicating if modality is present
+        batch_size = enc_inputs[0].size(0)
+        
+        for modality_encoders in self.fusion_stacks:
+            enc_outputs = []
+            bottleneck_tokens_modality_sum = torch.zeros_like(bottleneck_tokens)
+            sum_of_modalities = torch.zeros(batch_size, 1, 1, device=bottleneck_tokens.device)
+
+            for idx, enc_layer in enumerate(modality_encoders):
+                # Concatenate bottleneck tokens with modality tokens
+                # bottleneck_tokens: [B, num_bottlenecks, d_model]
+                # enc_inputs[idx]: [B, seq_len, d_model]
+                fused_input = torch.cat([bottleneck_tokens, enc_inputs[idx]], dim=1)
+                
+                # Padding mask for bottleneck tokens is always False (i.e. valid)
+                # [B, num_bottlenecks] of False
+                b_mask = torch.zeros(batch_size, self.n_bottlenecks, dtype=torch.bool, device=fused_input.device)
+                
+                # Modality padding mask
+                m_mask = ~masks[idx] if masks[idx] is not None else torch.zeros(batch_size, enc_inputs[idx].size(1), dtype=torch.bool, device=fused_input.device)
+                
+                combined_mask = torch.cat([b_mask, m_mask], dim=1)
+                
+                # Pass through the layer
+                enc_out = enc_layer(fused_input, src_key_padding_mask=combined_mask)
+                
+                # The output consists of processed bottleneck tokens and modality tokens
+                # [B, num_bottlenecks, d_model] and [B, seq_len, d_model]
+                bottleneck_hidden_tokens = enc_out[:, :self.n_bottlenecks, :]
+                modality_hidden_tokens = enc_out[:, self.n_bottlenecks:, :]
+                enc_outputs.append(modality_hidden_tokens)
+                
+                # Average updated bottlenecks from valid modalities
+                modality_is_valid = valid_modalities[idx].view(batch_size, 1, 1)
+                bottleneck_tokens_modality_sum += bottleneck_hidden_tokens * modality_is_valid
+                sum_of_modalities += modality_is_valid
+
+            # Prevent division by zero if all modalities are missing
+            # If sum_of_modalities is 0, just pass zeros (or keep previous bottleneck_tokens)
+            # sum_of_modalities = torch.clamp(sum_of_modalities, min=1.0)
+            avg_divisor = sum_of_modalities.clone()
+            avg_divisor[avg_divisor == 0] = 1.0
+
+            bottleneck_tokens = bottleneck_tokens_modality_sum / avg_divisor
+            enc_inputs = enc_outputs
+            
+        return enc_inputs
+
+    def forward(self, enc_inputs: List[torch.Tensor], masks: List[torch.Tensor]) -> List[torch.Tensor]:
+        batch_size = enc_inputs[0].size(0)
+
+        # Determine if a modality is valid for each instance in the batch
+        # A modality is valid if it has at least one True in its mask
+        valid_modalities = []
+        for mask, inp in zip(masks, enc_inputs):
+            if mask is not None:
+                # [B] - True if there's any valid token (1/True)
+                valid = mask.any(dim=1).float()
+            else:
+                valid = torch.ones(batch_size, device=inp.device)
+            valid_modalities.append(valid)
+            
+        bottleneck_tokens = self.bottlenecks.expand(batch_size, -1, -1)
+        
+        enc_inputs = self.forward_prefusion(enc_inputs, masks)
+        enc_inputs = self.forward_fusion(enc_inputs, masks, bottleneck_tokens, valid_modalities)
+        
+        return enc_inputs
+
+
+class BottleneckTransformer(BaseModel):
+    """Bottleneck Transformer model for PyHealth datasets.
+
+    Per-field mode: each feature stream is embedded with :class:`EmbeddingModel`,
+    prefixed with a learnable per-modality ``[CLS]`` token, processed by
+    independent prefusion layers, then fused via shared bottleneck tokens.
+    The per-modality ``[CLS]`` embeddings are averaged and fed to the
+    classification head.
+
+    Unified mode (``unified_embedding`` supplied): all temporal fields are
+    jointly embedded and time-sorted by
+    :class:`~pyhealth.models.embedding.unified.UnifiedMultimodalEmbeddingModel`
+    into a single interleaved sequence.  A single ``[CLS]`` token is prepended
+    and the encoder runs with ``n_modality=1``, so the bottleneck tokens attend
+    over the full cross-modal timeline.
+
+    Args:
+        dataset (SampleDataset): dataset providing processed inputs.
+        embedding_dim (int): shared embedding dimension.
+        bottlenecks_n (int): number of shared bottleneck tokens.
+        fusion_startidx (int): layer index at which bottleneck fusion starts.
+            Must satisfy ``0 <= fusion_startidx <= num_layers``.
+        num_layers (int): total transformer layers (prefusion + fusion).
+        heads (int): number of attention heads per transformer block.
+        dropout (float): dropout rate inside transformer blocks.
+        unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when
+            provided, switches to unified mode.
+
+    Examples:
+        >>> from pyhealth.datasets import create_sample_dataset, get_dataloader
+        >>> samples = [
+        ...     {
+        ...         "patient_id": "patient-0",
+        ...         "visit_id": "visit-0",
+        ...         "conditions": ["A", "B", "C"],
+        ...         "procedures": ["X", "Y"],
+        ...         "label": 1,
+        ...     },
+        ...     {
+        ...         "patient_id": "patient-1",
+        ...         "visit_id": "visit-0",
+        ...         "conditions": ["D"],
+        ...         "procedures": ["Z", "Y"],
+        ...         "label": 0,
+        ...     },
+        ... ]
+        >>> input_schema = {"conditions": "sequence", "procedures": "sequence"}
+        >>> output_schema = {"label": "binary"}
+        >>> dataset = create_sample_dataset(
+        ...     samples,
+        ...     input_schema,
+        ...     output_schema,
+        ...     dataset_name="demo",
+        ... )
+        >>> model = BottleneckTransformer(dataset=dataset, num_layers=3, fusion_startidx=1, bottlenecks_n=4)
+        >>> loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+        >>> batch = next(iter(loader))
+        >>> output = model(**batch)
+        >>> sorted(output.keys())
+        ['logit', 'loss', 'y_prob', 'y_true']
+    """
+
+    def __init__(
+        self,
+        dataset: SampleDataset,
+        embedding_dim: int = 128,
+        bottlenecks_n: int = 4,
+        fusion_startidx: int = 1,
+        num_layers: int = 3,
+        heads: int = 4,
+        dropout: float = 0.5,
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
+    ):
+        super().__init__(dataset=dataset)
+        self.embedding_dim = embedding_dim
+        self.bottlenecks_n = bottlenecks_n
+        self.fusion_startidx = fusion_startidx
+        self.num_layers = num_layers
+        self.heads = heads
+        self.dropout = dropout
+        self._use_unified = unified_embedding is not None
+
+        assert (
+            len(self.label_keys) == 1
+        ), "Only one label key is supported if BottleneckTransformer is initialized"
+        self.label_key = self.label_keys[0]
+        self.mode = self.dataset.output_schema[self.label_key]
+
+        assert 0 <= fusion_startidx <= num_layers, (
+            f"fusion_startidx must be in [0, num_layers], got {fusion_startidx}"
+        )
+
+        output_size = self.get_output_size()
+
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            # Single CLS token for the unified interleaved sequence
+            self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02)
+            self.encoder = MultimodalBottleneckTransformerEncoder(
+                n_modality=1,
+                bottlenecks_n=bottlenecks_n,
+                fusion_startidx=fusion_startidx,
+                n_layers=num_layers,
+                n_head=heads,
+                d_model=embedding_dim,
+                d_ff=embedding_dim * 4,
+                dropout=dropout,
+            )
+        else:
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            self.n_modality = len(self.feature_keys)
+            # Per-modality CLS tokens
+            self.cls_token_per_modality = nn.ParameterList([
+                nn.Parameter(torch.randn(1, 1, embedding_dim) * 0.02)
+                for _ in range(self.n_modality)
+            ])
+            self.encoder = MultimodalBottleneckTransformerEncoder(
+                n_modality=self.n_modality,
+                bottlenecks_n=bottlenecks_n,
+                fusion_startidx=fusion_startidx,
+                n_layers=num_layers,
+                n_head=heads,
+                d_model=embedding_dim,
+                d_ff=embedding_dim * 4,
+                dropout=dropout,
+            )
+
+        # fc input is embedding_dim in both modes (CLS token, not concat)
+        self.fc = nn.Linear(embedding_dim, output_size)
+
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Extract value/time/mask tensors for UnifiedMultimodalEmbeddingModel."""
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+        return inputs
+
+    def _forward_unified(
+        self,
+        **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+    ) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode.
+
+        Calls UnifiedMultimodalEmbeddingModel to produce a single time-sorted
+        sequence, prepends a CLS token, encodes with the bottleneck encoder
+        (n_modality=1), and classifies from the CLS output.
+        """
+        inputs = self._build_unified_inputs(kwargs)
+        out = self.embedding_model(inputs)
+        sequence = out["sequence"]           # (B, S, E)
+        event_mask = out["mask"].bool()      # (B, S)
+
+        # Prepend CLS token
+        batch_size = sequence.size(0)
+        cls = self.cls_token.expand(batch_size, -1, -1)
+        sequence = torch.cat([cls, sequence], dim=1)
+        cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=sequence.device)
+        event_mask = torch.cat([cls_mask, event_mask], dim=1)
+
+        enc_outputs = self.encoder([sequence], [event_mask])
+        patient_emb = enc_outputs[0][:, 0, :]   # CLS token output
+
+        logits = self.fc(patient_emb)
+        y_prob = self.prepare_y_prob(logits)
+        results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob}
+        if self.label_key in kwargs:
+            y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device)
+            results["loss"] = self.get_loss_function()(logits, y_true)
+            results["y_true"] = y_true
+        return results
+
+    @staticmethod
+    def _pool_embedding(x: torch.Tensor) -> torch.Tensor:
+        if x.dim() == 4:
+            x = x.sum(dim=2)
+        if x.dim() == 2:
+            x = x.unsqueeze(1)
+        return x
+
+    @staticmethod
+    def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor:
+        mask = torch.any(torch.abs(x) > 0, dim=-1)
+        if mask.dim() == 1:
+            mask = mask.unsqueeze(1)
+        invalid_rows = ~mask.any(dim=1)
+        if invalid_rows.any():
+            mask[invalid_rows, 0] = True
+        return mask.bool()
+
+    def forward(
+        self,
+        **kwargs: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
+    ) -> Dict[str, torch.Tensor]:
+        """Forward propagation.
+
+        In unified mode dispatches to :meth:`_forward_unified`. Otherwise runs
+        per-field embedding + bottleneck fusion.
+
+        Args:
+            **kwargs: keyword arguments for the model.
+
+        Returns:
+            A dictionary with the following keys:
+                loss: a scalar tensor representing the final loss.
+                y_prob: a tensor of predicted probabilities.
+                y_true: a tensor representing the true labels.
+                logit: the raw logits before activation.
+        """
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
+        enc_inputs = []
+        masks = []
+        
+        for idx, feature_key in enumerate(self.feature_keys):
+            feature = kwargs[feature_key]
+
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+
+            schema = self.dataset.input_processors[feature_key].schema()
+
+            value = feature[schema.index("value")] if "value" in schema else None
+            mask = feature[schema.index("mask")] if "mask" in schema else None
+
+            if len(feature) == len(schema) + 1 and mask is None:
+                mask = feature[-1]
+
+            if value is None:
+                raise ValueError(
+                    f"Feature '{feature_key}' must contain 'value' "
+                    f"in the schema."
+                )
+            else:
+                value = value.to(self.device)
+                
+            if mask is not None:
+                mask = mask.to(self.device)
+                value = self.embedding_model({feature_key: value}, masks={feature_key: mask})[feature_key]
+            else:
+                value = self.embedding_model({feature_key: value})[feature_key]
+                
+            value = self._pool_embedding(value)
+            
+            if mask is not None:
+                mask = mask.bool()
+                if mask.dim() == value.dim():
+                    mask = mask.any(dim=-1)
+            else:
+                mask = self._mask_from_embeddings(value)
+            
+            # Prepend Modality CLS token
+            batch_size = value.size(0)
+            cls_token = self.cls_token_per_modality[idx].expand(batch_size, -1, -1)
+            value = torch.cat([cls_token, value], dim=1)
+            
+            # Update mask for CLS token (always valid)
+            cls_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=value.device)
+            mask = torch.cat([cls_mask, mask], dim=1)
+            
+            enc_inputs.append(value)
+            masks.append(mask)
+
+        # Pass through Bottleneck Transformer Encoder
+        enc_outputs = self.encoder(enc_inputs, masks)
+        
+        # Extract CLS tokens
+        cls_tokens = [out[:, 0, :].unsqueeze(1) for out in enc_outputs]
+        cls_tokens = torch.cat(cls_tokens, dim=1) # [B, n_modality, embedding_dim]
+        
+        # Average CLS tokens across valid modalities
+        b_size = cls_tokens.size(0)
+        valid_modalities = []
+        for mask in masks:
+            # We check if there's any valid token aside from the CLS token (index 0)
+            if mask.size(1) > 1:
+                valid = mask[:, 1:].any(dim=1).float()
+            else:
+                valid = mask[:, 0].float() # fallback
+            valid_modalities.append(valid.view(b_size, 1, 1))
+            
+        valid_modality_tensor = torch.cat(valid_modalities, dim=1) # [B, n_modality, 1]
+        
+        # Apply valid mask
+        masked_cls = cls_tokens * valid_modality_tensor
+        sum_valid = valid_modality_tensor.sum(dim=1) # [B, 1]
+        
+        # Avoid division by zero
+        sum_valid[sum_valid == 0] = 1.0
+        patient_emb = masked_cls.sum(dim=1) / sum_valid # [B, embedding_dim]
+        
+        logits = self.fc(patient_emb)
+        y_prob = self.prepare_y_prob(logits)
+
+        results = {
+            "logit": logits,
+            "y_prob": y_prob,
+        }
+
+        if self.label_key in kwargs:
+            y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device)
+            loss = self.get_loss_function()(logits, y_true)
+            results["loss"] = loss
+            results["y_true"] = y_true
+
+        return results
+
+if __name__ == "__main__":
+    from pyhealth.datasets import create_sample_dataset, get_dataloader
+
+    samples = [
+        {
+            "patient_id": "patient-0",
+            "visit_id": "visit-0",
+            "conditions": ["A", "B", "C"],
+            "procedures": ["X", "Y"],
+            "label": 1,
+        },
+        {
+            "patient_id": "patient-1",
+            "visit_id": "visit-0",
+            "conditions": ["D"],
+            "procedures": ["Z", "Y"],
+            "label": 0,
+        },
+    ]
+
+    input_schema = {
+        "conditions": "sequence",
+        "procedures": "sequence",
+    }
+    output_schema = {"label": "binary"}
+
+    dataset = create_sample_dataset(
+        samples=samples,
+        input_schema=input_schema,
+        output_schema=output_schema,
+        dataset_name="test",
+    )
+
+    train_loader = get_dataloader(dataset, batch_size=2, shuffle=True)
+
+    model = BottleneckTransformer(
+        dataset=dataset, 
+        embedding_dim=64, 
+        bottlenecks_n=2, 
+        fusion_startidx=1, 
+        num_layers=3, 
+        heads=2
+    )
+
+    data_batch = next(iter(train_loader))
+
+    result = model(**data_batch)
+    print(result)
+
+    result["loss"].backward()
+    print("Test completed successfully.")
diff --git a/pyhealth/models/ehrmamba.py b/pyhealth/models/ehrmamba.py
index e24c5595f..4c20cc563 100644
--- a/pyhealth/models/ehrmamba.py
+++ b/pyhealth/models/ehrmamba.py
@@ -6,6 +6,7 @@
 from pyhealth.datasets import SampleDataset
 from pyhealth.models import BaseModel
 from pyhealth.models.embedding import EmbeddingModel
+from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel
 from pyhealth.models.utils import get_last_visit
 from pyhealth.processors import (
     MultiHotProcessor,
@@ -111,6 +112,11 @@ class EHRMamba(BaseModel):
     Electronic Health Records (arxiv 2405.14567). Uses Mamba (SSM) for linear
     complexity in sequence length; supports long EHR sequences.
 
+    When ``unified_embedding`` is supplied the model switches to **unified
+    mode**: all temporal fields are jointly embedded and time-sorted by
+    :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single*
+    stack of :class:`MambaBlock` layers rather than one stack per field.
+
     Args:
         dataset: SampleDataset for token/embedding setup.
         embedding_dim: Embedding and hidden dimension. Default 128.
@@ -118,6 +124,8 @@ class EHRMamba(BaseModel):
         state_size: SSM state size per channel. Default 16.
         conv_kernel: Causal conv kernel size in block. Default 4.
         dropout: Dropout before classification head. Default 0.1.
+        unified_embedding: Optional pre-built UnifiedMultimodalEmbeddingModel.
+            When provided, enables unified multi-modal mode.
     """
 
     def __init__(
@@ -128,6 +136,7 @@ def __init__(
         state_size: int = 16,
         conv_kernel: int = 4,
         dropout: float = 0.1,
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
     ):
         super().__init__(dataset=dataset)
         self.embedding_dim = embedding_dim
@@ -135,19 +144,18 @@ def __init__(
         self.state_size = state_size
         self.conv_kernel = conv_kernel
         self.dropout_rate = dropout
+        self._use_unified = unified_embedding is not None
 
         assert len(self.label_keys) == 1, "EHRMamba supports single label key only"
         self.label_key = self.label_keys[0]
         self.mode = self.dataset.output_schema[self.label_key]
 
-        self.embedding_model = EmbeddingModel(dataset, embedding_dim)
-        self.feature_processors = {
-            k: self.dataset.input_processors[k] for k in self.feature_keys
-        }
+        output_size = self.get_output_size()
+        self.dropout = nn.Dropout(dropout)
 
-        self.blocks = nn.ModuleDict()
-        for feature_key in self.feature_keys:
-            self.blocks[feature_key] = nn.ModuleList(
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            self._unified_blocks = nn.ModuleList(
                 [
                     MambaBlock(
                         d_model=embedding_dim,
@@ -157,10 +165,76 @@ def __init__(
                     for _ in range(num_layers)
                 ]
             )
-
-        output_size = self.get_output_size()
-        self.dropout = nn.Dropout(dropout)
-        self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size)
+            self.fc = nn.Linear(embedding_dim, output_size)
+        else:
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            self.feature_processors = {
+                k: self.dataset.input_processors[k] for k in self.feature_keys
+            }
+            self.blocks = nn.ModuleDict()
+            for feature_key in self.feature_keys:
+                self.blocks[feature_key] = nn.ModuleList(
+                    [
+                        MambaBlock(
+                            d_model=embedding_dim,
+                            state_size=state_size,
+                            conv_kernel=conv_kernel,
+                        )
+                        for _ in range(num_layers)
+                    ]
+                )
+            self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size)
+
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Build the inputs dict required by UnifiedMultimodalEmbeddingModel."""
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+        return inputs
+
+    def _forward_unified(self, **kwargs) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode.
+
+        Calls UnifiedMultimodalEmbeddingModel to produce a single
+        temporally-sorted event sequence, then encodes it with one shared
+        MambaBlock stack and pools to the last valid event.
+        """
+        inputs = self._build_unified_inputs(kwargs)
+        out = self.embedding_model(inputs)
+        x    = out["sequence"]       # (B, S_total, E)
+        mask = out["mask"].bool()    # (B, S_total)
+
+        for blk in self._unified_blocks:
+            x = blk(x)
+
+        last_h = get_last_visit(x, mask)
+        logits = self.fc(self.dropout(last_h))
+        y_prob = self.prepare_y_prob(logits)
+        results: Dict[str, torch.Tensor] = {
+            "loss": torch.tensor(0.0),  # placeholder, overwritten below
+            "y_prob": y_prob,
+            "logit": logits,
+        }
+        if self.label_key in kwargs:
+            y_true = kwargs[self.label_key].to(self.device)
+            results["loss"]   = self.get_loss_function()(logits, y_true)
+            results["y_true"] = y_true
+        if kwargs.get("embed", False):
+            results["embed"] = last_h
+        return results
 
     @staticmethod
     def _split_temporal(feature: Any) -> Tuple[Optional[torch.Tensor], Any]:
@@ -211,6 +285,9 @@ def _pool_embedding(x: torch.Tensor) -> torch.Tensor:
         return x
 
     def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
         patient_emb = []
         embedding_inputs: Dict[str, torch.Tensor] = {}
         masks: Dict[str, torch.Tensor] = {}
diff --git a/pyhealth/models/embedding/__init__.py b/pyhealth/models/embedding/__init__.py
new file mode 100644
index 000000000..8b0bd38b5
--- /dev/null
+++ b/pyhealth/models/embedding/__init__.py
@@ -0,0 +1,37 @@
+"""Embedding models for PyHealth multimodal pipelines.
+
+All embedding models share the :class:`BaseEmbeddingModel` interface:
+they expose an ``embedding_dim`` property and a ``forward`` method that
+transforms processor output tensors into dense vector embeddings.
+
+Available models:
+
+- :class:`EmbeddingModel` — generic encoder for codes, sequences, timeseries
+- :class:`VisionEmbeddingModel` — ViT-style patch encoder for medical images (Josh)
+- :class:`TextEmbeddingModel` — BERT-based encoder for clinical text (Rian)
+- :class:`UnifiedMultimodalEmbeddingModel` — temporally-aligned multi-modal encoder
+
+Helper utilities:
+
+- :class:`SinusoidalTimeEmbedding` — continuous time positional encoding
+- :func:`init_embedding_with_pretrained` — load GloVe-style pretrained vectors
+"""
+
+from .base import BaseEmbeddingModel
+from .vanilla import EmbeddingModel, init_embedding_with_pretrained
+from .vision import VisionEmbeddingModel, PatchEmbedding, Permute
+from .text import TextEmbeddingModel, TextEmbedding
+from .unified import UnifiedMultimodalEmbeddingModel, SinusoidalTimeEmbedding
+
+__all__ = [
+    "BaseEmbeddingModel",
+    "EmbeddingModel",
+    "VisionEmbeddingModel",
+    "PatchEmbedding",
+    "Permute",
+    "TextEmbeddingModel",
+    "TextEmbedding",
+    "UnifiedMultimodalEmbeddingModel",
+    "SinusoidalTimeEmbedding",
+    "init_embedding_with_pretrained",
+]
diff --git a/pyhealth/models/embedding/base.py b/pyhealth/models/embedding/base.py
new file mode 100644
index 000000000..5ff1ab2a6
--- /dev/null
+++ b/pyhealth/models/embedding/base.py
@@ -0,0 +1,52 @@
+from abc import ABC, abstractmethod
+
+
+class BaseEmbeddingModel(ABC):
+    """Abstract base class for all embedding models in PyHealth.
+
+    All embedding models share a common contract:
+
+    - They expose an ``embedding_dim`` property indicating the output vector dimension.
+    - Their ``forward`` method accepts processor output tensors and returns
+      vector embeddings.
+
+    Concrete subclasses:
+
+    - :class:`EmbeddingModel` – generic encoder for codes, sequences, timeseries
+    - :class:`VisionEmbeddingModel` – patch-based encoder for medical images (Josh)
+    - :class:`TextEmbeddingModel` – BERT-based encoder for clinical text (Rian)
+    - :class:`UnifiedMultimodalEmbeddingModel` – temporally-aligned multi-modal encoder
+    """
+
+    @property
+    @abstractmethod
+    def embedding_dim(self) -> int:
+        """Output embedding dimension shared across all modalities."""
+        ...
+
+    @abstractmethod
+    def forward(self, *args, **kwargs):
+        """Transform processor outputs into embeddings.
+
+        Subclass return types
+        ---------------------
+        EmbeddingModel
+            ``Dict[str, Tensor]`` mapping each field name to its embedded
+            tensor.  When ``output_mask=True`` is passed, returns a
+            ``(Dict[str, Tensor], Dict[str, Tensor])`` tuple of
+            (embeddings, masks).
+
+        VisionEmbeddingModel
+            ``Tensor`` of shape ``[batch, embedding_dim]``.
+
+        TextEmbeddingModel
+            ``(Tensor, BoolTensor)`` of shapes ``([B, T, E], [B, T])``
+            when ``return_mask=True`` (default), or a plain
+            ``Tensor [B, T, E]`` when ``return_mask=False``.
+
+        UnifiedMultimodalEmbeddingModel
+            ``Dict[str, Tensor]`` with keys ``"sequence"`` ``[B, S, E]``,
+            ``"mask"`` ``[B, S]``, ``"time"`` ``[B, S]``, and
+            ``"type_ids"`` ``[B, S]``.
+        """
+        ...
diff --git a/pyhealth/models/text_embedding.py b/pyhealth/models/embedding/text.py
similarity index 91%
rename from pyhealth/models/text_embedding.py
rename to pyhealth/models/embedding/text.py
index 57549fea5..87aa11ea0 100644
--- a/pyhealth/models/text_embedding.py
+++ b/pyhealth/models/embedding/text.py
@@ -1,5 +1,7 @@
 """Text embedding module for multimodal PyHealth pipelines.
 
+Author: Rian
+
 This module provides a Transformer-based text encoder for clinical/medical text.
 It is designed to integrate with PyHealth's multimodal fusion architecture.
 
@@ -14,8 +16,8 @@
     - torch
 
 Example:
-    >>> from pyhealth.models.text_embedding import TextEmbedding
-    >>> encoder = TextEmbedding(embedding_dim=256)
+    >>> from pyhealth.models.embedding import TextEmbeddingModel
+    >>> encoder = TextEmbeddingModel(embedding_dim=256)
     >>> embeddings, mask = encoder(["Patient has fever.", "Follow-up."])
     >>> embeddings.shape  # [2, T, 256]
 """
@@ -28,11 +30,13 @@
 import torch.nn as nn
 from transformers import AutoModel, AutoTokenizer
 
+from .base import BaseEmbeddingModel
+
 
 logger = logging.getLogger(__name__)
 
 
-class TextEmbedding(nn.Module):
+class TextEmbeddingModel(nn.Module, BaseEmbeddingModel):
     """Encodes clinical text into embeddings for multimodal fusion.
 
     This module wraps a pretrained Hugging Face transformer (default:
@@ -55,7 +59,7 @@ class TextEmbedding(nn.Module):
 
         Example: A 300-token note with chunk_size=128 becomes 3 chunks:
             Chunk 1: [CLS] + tokens[0:126] + [SEP]   = 128 tokens
-            Chunk 2: [CLS] + tokens[126:252] + [SEP] = 128 tokens  
+            Chunk 2: [CLS] + tokens[126:252] + [SEP] = 128 tokens
             Chunk 3: [CLS] + tokens[252:300] + [SEP] = 50 tokens
 
     Pooling Modes:
@@ -117,7 +121,7 @@ class TextEmbedding(nn.Module):
     Example:
         Basic usage with default parameters:
 
-        >>> encoder = TextEmbedding(embedding_dim=256)
+        >>> encoder = TextEmbeddingModel(embedding_dim=256)
         >>> texts = ["Patient presents with chest pain.", "Routine checkup."]
         >>> embeddings, mask = encoder(texts)
         >>> embeddings.shape
@@ -127,14 +131,14 @@ class TextEmbedding(nn.Module):
 
         Using chunk-level pooling for efficiency:
 
-        >>> encoder = TextEmbedding(pooling="cls", embedding_dim=128)
+        >>> encoder = TextEmbeddingModel(pooling="cls", embedding_dim=128)
         >>> long_note = "..." * 1000  # Very long clinical note
         >>> emb, mask = encoder([long_note])
         >>> emb.shape  # [1, num_chunks, 128] instead of [1, thousands, 128]
 
         Backward-compatible single tensor return:
 
-        >>> encoder = TextEmbedding(return_mask=False)
+        >>> encoder = TextEmbeddingModel(return_mask=False)
         >>> embeddings = encoder(["Test"])  # Just tensor, no tuple
     """
 
@@ -159,7 +163,7 @@ def __init__(
         """
         super().__init__()
         self.model_name = model_name
-        self.embedding_dim = embedding_dim
+        self._embedding_dim = embedding_dim
         self.chunk_size = chunk_size
         self.max_chunks = max_chunks
         self.pooling = pooling
@@ -184,6 +188,10 @@ def __init__(
         # This aligns text embeddings with other modalities in a shared E' space
         self.fc = nn.Linear(self.transformer.config.hidden_size, embedding_dim)
 
+    @property
+    def embedding_dim(self) -> int:
+        return self._embedding_dim
+
     def _chunk_and_encode(
         self, text: str, device: torch.device
     ) -> torch.Tensor:
@@ -232,11 +240,6 @@ def _chunk_and_encode(
             chunks = [[self.tokenizer.cls_token_id, self.tokenizer.sep_token_id]]
 
         # Step 3: Apply max_chunks limit (performance guardrail)
-        # Rationale: Clinical notes can be 10K+ tokens. Without a cap:
-        # - Memory usage explodes (each chunk needs transformer forward pass)
-        # - Silent OOMs in production environments
-        # - Inference time becomes unpredictable
-        # We warn rather than silently truncate so users can adjust.
         if self.max_chunks is not None and len(chunks) > self.max_chunks:
             original_chunks = len(chunks)
             chunks = chunks[: self.max_chunks]
@@ -319,18 +322,6 @@ def forward(
 
             If return_mask=False (backward compatibility):
                 torch.Tensor: Just the embeddings tensor [B, T, E']
-
-        Note:
-            The return_mask parameter exists for backward compatibility.
-            New code should use the default return_mask=True to get masks
-            needed for downstream attention layers.
-
-        Example:
-            >>> encoder = TextEmbedding(embedding_dim=128)
-            >>> emb, mask = encoder(["Hello world", "A longer text here"])
-            >>> emb.shape   # [2, T, 128] where T is max tokens
-            >>> mask.shape  # [2, T]
-            >>> mask[0].sum()  # Number of valid tokens in first sample
         """
         # Normalize single string to list
         if isinstance(text, str):
@@ -357,7 +348,7 @@ def forward(
 
             # Pad embedding tensor with zeros
             if pad_len > 0:
-                padding = torch.zeros(pad_len, self.embedding_dim, device=device)
+                padding = torch.zeros(pad_len, self._embedding_dim, device=device)
                 e = torch.cat([e, padding], dim=0)
             padded.append(e)
 
@@ -377,3 +368,7 @@ def forward(
             return embeddings, mask
         else:
             return embeddings
+
+
+# Alias for backward compatibility
+TextEmbedding = TextEmbeddingModel
diff --git a/pyhealth/models/embedding/unified.py b/pyhealth/models/embedding/unified.py
new file mode 100644
index 000000000..78bc02cfc
--- /dev/null
+++ b/pyhealth/models/embedding/unified.py
@@ -0,0 +1,595 @@
+"""UnifiedMultimodalEmbeddingModel, temporally aligned multimodal embedding.
+
+Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor``
+subclasses ), embeds each event with a modality-specific encoder, then
+interleaves all events on a shared timeline by sorting on timestamp and adding
+sinusoidal time embeddings + learned modality-type embeddings.
+
+Output shape: ``(B, S_total, E')``, a single sequence of events usable by
+any downstream sequence model (Transformer, Mamba, RNN, …).
+
+IMAGE encoding delegates to :class:`PatchEmbedding` from
+:mod:`pyhealth.models.embedding.vision` (Josh's model), pooling patch tokens
+to a single per-image vector via global mean pooling.
+
+TEXT encoding uses a pretrained BERT tokenizer model directly, extracting the
+[CLS] token per note, the same BERT-based approach as
+:class:`TextEmbeddingModel` (Rian's model).
+
+Unimodal model reuse via ``field_embeddings``::
+
+    vision_model = VisionEmbeddingModel(dataset, embedding_dim=128)
+    text_model   = TextEmbeddingModel(embedding_dim=128)
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=128,
+        field_embeddings={
+            "chest_xray": vision_model,   # reuses trained backbone
+            "notes":      text_model,     # reuses BERT + projection
+        },
+    )
+
+Quickstart::
+
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.datasets.collate import collate_temporal
+    model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128)
+    # inside forward:
+    #   inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...}
+    out = model(inputs)
+    # out["sequence"]: (B, S_total, 128)
+    # out["mask"]:     (B, S_total)     , 1 = real event, 0 = padding
+    # out["time"]:     (B, S_total)     , hours from first event
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+
+from ...processors.base_processor import ModalityType, TemporalFeatureProcessor
+from .base import BaseEmbeddingModel
+from .vision import PatchEmbedding
+
+
+# ── Helpers ───────────────────────────────────────────────────────────────────
+
+
+class SinusoidalTimeEmbedding(nn.Module):
+    """Continuous sinusoidal embedding for scalar time values (in hours).
+
+    Identical in spirit to the positional encoding in "Attention is All You
+    Need" but operating on real-valued timestamps rather than integer positions.
+
+    Args:
+        dim: Output embedding dimension (must be even).
+        max_hours: Maximum expected time value in hours.  Values are normalised
+            to ``[0, 2π]`` before the sin/cos projection.  Default 720 (30 days).
+
+    Shape:
+        Input:  ``(*, )``  float tensor of times in hours
+        Output: ``(*, dim)``
+    """
+
+    def __init__(self, dim: int, max_hours: float = 720.0):
+        super().__init__()
+        assert dim % 2 == 0, f"dim must be even, got {dim}"
+        self.dim = dim
+        self.max_hours = max_hours
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1)
+        )
+        self.register_buffer("freqs", freqs)  # (dim//2,)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        """:param t: ``(...,)`` float, times in hours."""
+        t_norm = t / self.max_hours * 2 * math.pi  # (...,)
+        args = t_norm.unsqueeze(-1) * self.freqs  # (..., dim//2)
+        return torch.cat([args.sin(), args.cos()], dim=-1)  # (..., dim)
+
+
+class _MeanPool(nn.Module):
+    """Pool a sequence of patch embeddings to a single vector via global mean."""
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # (B, num_patches, E) -> (B, E)
+        return x.mean(dim=1)
+
+
+# ── Main model ───────────────────────────────────────────────────────────────
+
+
+class UnifiedMultimodalEmbeddingModel(nn.Module, BaseEmbeddingModel):
+    """Embed heterogeneous temporal features into a single aligned sequence.
+
+    **All** input processors must be ``TemporalFeatureProcessor`` subclasses.
+    Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``)
+    are rejected with a clear error, use :class:`EmbeddingModel` for those fields.
+
+    Modality routing:
+
+    - **CODE**: ``nn.Embedding`` lookup.
+    - **TEXT**: Pretrained BERT (same approach as :class:`TextEmbeddingModel`),
+      CLS token extracted per note.
+    - **IMAGE**: :class:`PatchEmbedding` (from :class:`VisionEmbeddingModel`)
+      followed by global mean pooling to produce one vector per image event.
+    - **NUMERIC / SIGNAL**: ``nn.Linear`` projection.
+
+    Unimodal model reuse:
+
+    Pass pre-built :class:`EmbeddingModel`, :class:`VisionEmbeddingModel`, or
+    :class:`TextEmbeddingModel` instances via ``field_embeddings`` to reuse
+    their trained encoder weights instead of building new ones from scratch.
+    The core encoder module is extracted from each pre-built model:
+
+    - ``EmbeddingModel`` → ``embedding_layers[field_name]`` (``nn.Embedding`` /
+      ``nn.Linear``)
+    - ``VisionEmbeddingModel`` → ``embedding_layers[field_name]`` backbone +
+      global mean pooling
+    - ``TextEmbeddingModel`` → ``transformer`` (BERT) + ``fc`` (projection)
+
+    Algorithm
+    ---------
+    For each temporal field:
+
+    1. Route ``inputs[field]["value"]`` through a modality-specific encoder →
+       ``(B, N_i, E')`` per-event embeddings.
+    2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours).
+    3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or
+       ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if
+       token-level.
+
+    Then:
+
+    4. Concatenate across all fields → ``(B, S_total, E')``.
+    5. Sort events along dim=1 by timestamp (ascending).
+    6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``.
+    7. Return ``{"sequence", "time", "mask", "type_ids"}``.
+
+    Args:
+        processors: ``dict[field_name, TemporalFeatureProcessor]``, the
+            processors for each temporal field in the dataset.  Pass
+            ``dataset.input_processors`` directly.
+        embedding_dim: Shared embedding dimension ``E'``.
+        time_embedding: ``"sinusoidal"`` (default) or ``"learned"``.
+        max_time_hours: Normalisation constant for the time embedding.
+            Defaults to 720 h (30 days).
+        image_size: Image size (H=W) assumed for IMAGE fields when using
+            PatchEmbedding. Defaults to 224.
+        image_channels: Number of input channels for IMAGE fields. Defaults to 3.
+        patch_size: Patch size for IMAGE PatchEmbedding encoder. Defaults to 16.
+        image_pool: Pooling strategy applied to IMAGE patch tokens to produce
+            one vector per image event. Only ``"mean"`` (global mean pooling)
+            is currently implemented. Defaults to ``"mean"``.
+        field_embeddings: Optional mapping of field names to pre-built unimodal
+            embedding models.  Supported types:
+
+            - :class:`EmbeddingModel` (codes / numeric) — extracts
+              ``embedding_layers[field_name]``.
+            - :class:`VisionEmbeddingModel` — extracts the backbone layer and
+              wraps it with global mean pooling.
+            - :class:`TextEmbeddingModel` — reuses ``transformer`` and ``fc``
+              for BERT-based CLS extraction.
+
+            Fields not present in this dict fall back to the default
+            internally-built encoders.
+
+    Example::
+
+        model = UnifiedMultimodalEmbeddingModel(
+            processors=dataset.input_processors,
+            embedding_dim=128,
+        )
+        # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}}
+        out = model(inputs)
+        seq = out["sequence"]   # (B, S_total, 128)
+        mask = out["mask"]      # (B, S_total)  float, 1=valid 0=pad
+
+        # With pre-built unimodal models:
+        vision = VisionEmbeddingModel(dataset, embedding_dim=128)
+        model = UnifiedMultimodalEmbeddingModel(
+            processors=dataset.input_processors,
+            embedding_dim=128,
+            field_embeddings={"chest_xray": vision},
+        )
+    """
+
+    def __init__(
+        self,
+        processors: dict[str, Any],
+        embedding_dim: int = 128,
+        time_embedding: str = "sinusoidal",
+        max_time_hours: float = 720.0,
+        image_size: int = 224,
+        image_channels: int = 3,
+        patch_size: int = 16,
+        image_pool: str = "mean",
+        field_embeddings: Optional[dict[str, Any]] = None,
+        freeze_text_encoder: bool = False,
+    ):
+        super().__init__()
+        if image_pool != "mean":
+            raise NotImplementedError(
+                f"Only image_pool='mean' is implemented, got {image_pool!r}."
+            )
+        self._embedding_dim = embedding_dim
+        self._freeze_text_encoder = freeze_text_encoder
+        self.image_pool = image_pool
+        _field_embeddings = field_embeddings or {}
+
+        self.encoders: nn.ModuleDict = nn.ModuleDict()
+        self.projections: nn.ModuleDict = nn.ModuleDict()
+        self.modality_types: dict[str, ModalityType] = {}
+        self._shared_text_field_by_model: dict[str, str] = {}
+        self._text_canonical: dict[str, str] = {}  # field → first field sharing the same tokenizer
+
+        for field_name, processor in processors.items():
+            if not isinstance(processor, TemporalFeatureProcessor):
+                raise TypeError(
+                    f"UnifiedMultimodalEmbeddingModel requires every input processor "
+                    f"to be a TemporalFeatureProcessor subclass, but '{field_name}' "
+                    f"uses {type(processor).__name__}.  For non-temporal fields use "
+                    f"EmbeddingModel."
+                )
+
+            m = processor.modality()
+            self.modality_types[field_name] = m
+            pre_built = _field_embeddings.get(field_name)
+
+            if m == ModalityType.CODE:
+                self.encoders[field_name] = self._build_code_encoder(
+                    field_name, processor, pre_built, embedding_dim
+                )
+
+            elif m == ModalityType.TEXT:
+                self._build_text_encoder(
+                    field_name, processor, pre_built, embedding_dim,
+                    freeze=freeze_text_encoder,
+                )
+
+            elif m == ModalityType.IMAGE:
+                self.encoders[field_name] = self._build_image_encoder(
+                    field_name,
+                    processor,
+                    pre_built,
+                    embedding_dim,
+                    image_size,
+                    image_channels,
+                    patch_size,
+                    image_pool,
+                )
+
+            elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL):
+                self.encoders[field_name] = self._build_numeric_encoder(
+                    field_name, processor, pre_built, embedding_dim
+                )
+
+            else:
+                raise NotImplementedError(
+                    f"No encoder implemented for modality {m!r} (field '{field_name}')."
+                )
+
+        # Shared type embedding, one vector per unique modality in this dataset
+        unique_modalities = sorted(set(self.modality_types.values()))
+        self._modality_to_idx: dict[ModalityType, int] = {
+            mod: i for i, mod in enumerate(unique_modalities)
+        }
+        self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim)
+        self._warned_nested_code_flatten = False
+
+        # Time embedding
+        if time_embedding == "sinusoidal":
+            self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours)
+        else:
+            raise NotImplementedError(
+                "Only 'sinusoidal' time embedding is implemented."
+            )
+
+    # ── Encoder builders ──────────────────────────────────────────────────────
+
+    def _build_code_encoder(
+        self,
+        field_name: str,
+        processor: TemporalFeatureProcessor,
+        pre_built: Any,
+        embedding_dim: int,
+    ) -> nn.Module:
+        """Build CODE encoder: nn.Embedding, optionally from a pre-built EmbeddingModel."""
+        if (
+            pre_built is not None
+            and hasattr(pre_built, "embedding_layers")
+            and field_name in pre_built.embedding_layers
+        ):
+            layer = pre_built.embedding_layers[field_name]
+            pre_dim = getattr(pre_built, "embedding_dim", embedding_dim)
+            if pre_dim != embedding_dim:
+                return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim))
+            return layer
+
+        vocab_size = processor.value_dim()
+        return nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
+
+    def _build_text_encoder(
+        self,
+        field_name: str,
+        processor: TemporalFeatureProcessor,
+        pre_built: Any,
+        embedding_dim: int,
+        freeze: bool = False,
+    ) -> None:
+        """Build TEXT encoder: BERT + projection, optionally from TextEmbeddingModel."""
+
+        def _set_projection(
+            pre_dim: int, proj_source: Optional[nn.Module] = None
+        ) -> None:
+            if pre_dim != embedding_dim:
+                if proj_source is not None:
+                    self.projections[field_name] = nn.Sequential(
+                        proj_source,
+                        nn.Linear(pre_dim, embedding_dim),
+                    )
+                else:
+                    self.projections[field_name] = nn.Linear(pre_dim, embedding_dim)
+            elif proj_source is not None:
+                self.projections[field_name] = proj_source
+
+        if (
+            pre_built is not None
+            and hasattr(pre_built, "transformer")
+            and hasattr(pre_built, "fc")
+        ):
+            self.encoders[field_name] = pre_built.transformer
+            if freeze:
+                for p in pre_built.transformer.parameters():
+                    p.requires_grad = False
+            pre_dim = getattr(pre_built, "embedding_dim", embedding_dim)
+            _set_projection(pre_dim, pre_built.fc)
+            return
+
+        if processor.is_token():
+            from transformers import AutoModel
+
+            bert = AutoModel.from_pretrained(processor.tokenizer_model)
+            if freeze:
+                for p in bert.parameters():
+                    p.requires_grad = False
+            self.encoders[field_name] = bert
+            hidden = bert.config.hidden_size
+            if hidden != embedding_dim:
+                self.projections[field_name] = nn.Linear(hidden, embedding_dim)
+        else:
+            raise ValueError(
+                f"TEXT processor '{field_name}' must either supply a pre-built "
+                f"TextEmbeddingModel via field_embeddings or use a tokenizer "
+                f"(set tokenizer_model=...) to be used with "
+                f"UnifiedMultimodalEmbeddingModel."
+            )
+
+    def _build_image_encoder(
+        self,
+        field_name: str,
+        processor: TemporalFeatureProcessor,
+        pre_built: Any,
+        embedding_dim: int,
+        image_size: int,
+        image_channels: int,
+        patch_size: int,
+        image_pool: str,
+    ) -> nn.Module:
+        """Build IMAGE encoder: backbone + pool, optionally from VisionEmbeddingModel."""
+        pool_layers: dict[str, nn.Module] = {"mean": _MeanPool()}
+        pool_layer = pool_layers[image_pool]
+
+        if (
+            pre_built is not None
+            and hasattr(pre_built, "embedding_layers")
+            and field_name in pre_built.embedding_layers
+        ):
+            backbone = pre_built.embedding_layers[field_name]
+            pre_dim = getattr(pre_built, "embedding_dim", embedding_dim)
+            if pre_dim != embedding_dim:
+                return nn.Sequential(
+                    backbone, pool_layer, nn.Linear(pre_dim, embedding_dim)
+                )
+            return nn.Sequential(backbone, pool_layer)
+
+        _image_size = getattr(processor, "image_size", image_size)
+        _in_channels = getattr(processor, "in_channels", image_channels)
+        return nn.Sequential(
+            PatchEmbedding(_image_size, patch_size, _in_channels, embedding_dim),
+            pool_layer,
+        )
+
+    def _build_numeric_encoder(
+        self,
+        field_name: str,
+        processor: TemporalFeatureProcessor,
+        pre_built: Any,
+        embedding_dim: int,
+    ) -> nn.Module:
+        """Build NUMERIC/SIGNAL encoder: nn.Linear, optionally from EmbeddingModel."""
+        if (
+            pre_built is not None
+            and hasattr(pre_built, "embedding_layers")
+            and field_name in pre_built.embedding_layers
+        ):
+            layer = pre_built.embedding_layers[field_name]
+            pre_dim = getattr(pre_built, "embedding_dim", embedding_dim)
+            if pre_dim != embedding_dim:
+                return nn.Sequential(layer, nn.Linear(pre_dim, embedding_dim))
+            return layer
+
+        in_features = processor.value_dim()
+        return nn.Linear(in_features, embedding_dim)
+
+    @property
+    def embedding_dim(self) -> int:
+        return self._embedding_dim
+
+    # ── Forward ───────────────────────────────────────────────────────────────
+
+    def forward(
+        self,
+        inputs: dict[str, dict[str, torch.Tensor]],
+    ) -> dict[str, torch.Tensor]:
+        """Encode and temporally align all temporal features.
+
+        Args:
+            inputs: ``{field_name: {"value": Tensor, "time": Tensor,
+                        "mask": Tensor (optional)}}``
+               , one dict per temporal feature, exactly as produced by
+                  ``collate_temporal``.
+
+        Returns:
+            A dict with keys:
+
+            * ``"sequence"``, ``(B, S_total, E')``  temporally-sorted events
+              (content + time + type embeddings)
+            * ``"time"``    , ``(B, S_total)``       timestamps (hours)
+            * ``"mask"``    , ``(B, S_total)``       1=real event, 0=padding
+            * ``"type_ids"``, ``(B, S_total)``       modality index per event
+            * ``"token_emb"``, ``(B, S_total, E')`` content-only event embedding
+              (before time/type are added); the target for masked modeling.
+        """
+        all_embeddings: list[torch.Tensor] = []
+        all_times: list[torch.Tensor] = []
+        all_masks: list[torch.Tensor] = []
+        all_types: list[torch.Tensor] = []
+
+        for field_name, feat_dict in inputs.items():
+            value = feat_dict["value"]  # (B, N_i, ...) or (B, S, F)
+            time = feat_dict["time"]  # (B, N_i)
+            mask = feat_dict.get("mask")
+
+            if time is None:
+                # Fallback: treat every event as occurring at t=0
+                time = torch.zeros(value.shape[:2], device=value.device)
+
+            modality = self.modality_types[field_name]
+            encoder_key = self._text_canonical.get(field_name, field_name)
+            encoder = self.encoders[encoder_key]
+
+            # ── Encode ────────────────────────────────────────────────────
+            if modality == ModalityType.CODE:
+                # CODE values may be either:
+                # - flat indices: (B, S)
+                # - nested indices: (B, S, C) where C is codes-per-event
+                # For nested indices, flatten to (B, S*C, E') so code-level
+                # detail is preserved, and expand time/mask to match.
+                if value.dim() == 2:
+                    emb = encoder(value)  # (B, S, E')
+                elif value.dim() == 3:
+                    bsz, seq_len, per_event_codes = value.shape
+                    token_emb = encoder(value.long())  # (B, S, C, E')
+                    emb = token_emb.reshape(bsz, seq_len * per_event_codes, -1)
+
+                    if not self._warned_nested_code_flatten:
+                        warnings.warn(
+                            (
+                                "UnifiedMultimodalEmbeddingModel detected "
+                                f"nested CODE input for '{field_name}' with "
+                                f"shape={tuple(value.shape)}. Flattening to "
+                                f"(B, S*C, E) and repeating time along C."
+                            ),
+                            stacklevel=2,
+                        )
+                        self._warned_nested_code_flatten = True
+
+                    if time is not None:
+                        time = (
+                            time.unsqueeze(-1)
+                            .expand(-1, -1, per_event_codes)
+                            .reshape(bsz, seq_len * per_event_codes)
+                        )
+
+                    if mask is not None:
+                        if mask.dim() == 2:
+                            mask = (
+                                mask.unsqueeze(-1)
+                                .expand(-1, -1, per_event_codes)
+                                .reshape(bsz, seq_len * per_event_codes)
+                            )
+                        elif mask.dim() == 3:
+                            mask = mask.reshape(bsz, seq_len * per_event_codes)
+                else:
+                    raise ValueError(
+                        f"Unsupported CODE value rank for '{field_name}': "
+                        f"shape={tuple(value.shape)}"
+                    )
+
+            elif modality == ModalityType.TEXT:
+                b, n, l = value.shape
+                flat_ids = value.view(b * n, l)
+                flat_mask = mask.view(b * n, l) if mask is not None else None
+                out = encoder(input_ids=flat_ids, attention_mask=flat_mask)
+                cls_emb = out.last_hidden_state[:, 0, :]  # (B*N, H)
+                if field_name in self.projections:
+                    cls_emb = self.projections[field_name](cls_emb)
+                emb = cls_emb.view(b, n, -1)  # (B, N, E')
+
+            elif modality == ModalityType.IMAGE:
+                # encoder = Sequential(PatchEmbedding, _MeanPool) → (B*N, E')
+                b, n, c, h, w = value.shape
+                flat_imgs = value.view(b * n, c, h, w)
+                img_emb = encoder(flat_imgs)  # (B*N, E')
+                emb = img_emb.view(b, n, -1)  # (B, N, E')
+
+            else:  # NUMERIC / SIGNAL
+                emb = encoder(value)  # (B, T, E')
+
+            # ── Build event-level validity mask ───────────────────────────
+            if mask is None:
+                event_mask = torch.ones(emb.shape[:2], device=emb.device)
+            else:
+                if mask.dim() > time.dim():
+                    # token-level (B, N, L) → event-level (B, N)
+                    event_mask = (mask.sum(dim=-1) > 0).float()
+                else:
+                    event_mask = mask.float()
+
+            # ── Modality type indices ─────────────────────────────────────
+            type_idx = self._modality_to_idx[modality]
+            type_ids = torch.full(
+                emb.shape[:2], type_idx, dtype=torch.long, device=emb.device
+            )
+
+            all_embeddings.append(emb)
+            all_times.append(time)
+            all_masks.append(event_mask)
+            all_types.append(type_ids)
+
+        # ── Concatenate across all fields ─────────────────────────────────
+        cat_emb = torch.cat(all_embeddings, dim=1)  # (B, S_total, E')
+        cat_time = torch.cat(all_times, dim=1)  # (B, S_total)
+        cat_mask = torch.cat(all_masks, dim=1)  # (B, S_total)
+        cat_types = torch.cat(all_types, dim=1)  # (B, S_total)
+
+        # ── Sort by time ──────────────────────────────────────────────────
+        sort_idx = cat_time.argsort(dim=1)
+        cat_emb = cat_emb.gather(1, sort_idx.unsqueeze(-1).expand_as(cat_emb))
+        cat_time = cat_time.gather(1, sort_idx)
+        cat_mask = cat_mask.gather(1, sort_idx)
+        cat_types = cat_types.gather(1, sort_idx)
+
+        # ── Add time + type embeddings ────────────────────────────────────
+        time_emb = self.time_embed(cat_time)  # (B, S_total, E')
+        type_emb = self.type_embedding(cat_types)  # (B, S_total, E')
+        final = cat_emb + time_emb + type_emb  # (B, S_total, E')
+
+        return {
+            "sequence": final,  # (B, S_total, E')
+            "time": cat_time,  # (B, S_total)
+            "mask": cat_mask,  # (B, S_total)
+            "type_ids": cat_types,  # (B, S_total)
+            # Per-event content embedding BEFORE time/type are added (same sort
+            # order as ``sequence``).  Masked-modeling pretrainers should
+            # reconstruct THIS rather than ``sequence``: the time/type
+            # components are largely recoverable from event position, so
+            # including them in the target dilutes the content signal.
+            "token_emb": cat_emb,   # (B, S_total, E')
+        }
diff --git a/pyhealth/models/embedding.py b/pyhealth/models/embedding/vanilla.py
similarity index 97%
rename from pyhealth/models/embedding.py
rename to pyhealth/models/embedding/vanilla.py
index 83a3a78c0..08a90f532 100644
--- a/pyhealth/models/embedding.py
+++ b/pyhealth/models/embedding/vanilla.py
@@ -6,8 +6,8 @@
 import torch
 import torch.nn as nn
 
-from ..datasets import SampleDataset
-from ..processors import (
+from ...datasets import SampleDataset
+from ...processors import (
     MultiHotProcessor,
     NestedFloatsProcessor,
     NestedSequenceProcessor,
@@ -19,7 +19,8 @@
     DeepNestedSequenceProcessor,
     DeepNestedFloatsProcessor,
 )
-from .base_model import BaseModel
+from ..base_model import BaseModel
+from .base import BaseEmbeddingModel
 
 
 def _iter_text_vectors(
@@ -147,7 +148,10 @@ def __init__(
         normalize_pretrained: bool = False,
     ):
         super().__init__(dataset)
-        self.embedding_dim = embedding_dim
+        # BaseEmbeddingModel declares `embedding_dim` as an abstract property,
+        # so we can't set self.embedding_dim directly (no setter).  Use a
+        # private backing attribute and expose it through the property below.
+        self._embedding_dim = embedding_dim
         self.embedding_layers = nn.ModuleDict()
 
         for field_name, processor in self.dataset.input_processors.items():
@@ -340,7 +344,6 @@ def forward(
 
             if output_mask:
                 # Generate a mask for this field
-                # For transformers, we might already have a mask, or use pad token
                 if masks is not None and field_name in masks:
                     out_masks[field_name] = masks[field_name].to(self.device)
                 elif hasattr(processor, "code_vocab"):
diff --git a/pyhealth/models/vision_embedding.py b/pyhealth/models/embedding/vision.py
similarity index 88%
rename from pyhealth/models/vision_embedding.py
rename to pyhealth/models/embedding/vision.py
index 228b87ddc..57eedda22 100644
--- a/pyhealth/models/vision_embedding.py
+++ b/pyhealth/models/embedding/vision.py
@@ -6,9 +6,10 @@
 import torch
 import torch.nn as nn
 import shutil
-from pyhealth.datasets import SampleDataset
-from pyhealth.models.base_model import BaseModel
-from pyhealth.processors import ImageProcessor
+from ...datasets import SampleDataset
+from ..base_model import BaseModel
+from ...processors import ImageProcessor
+from .base import BaseEmbeddingModel
 
 
 class Permute(nn.Module):
@@ -76,7 +77,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x
 
 
-class VisionEmbeddingModel(BaseModel):
+class VisionEmbeddingModel(BaseModel, BaseEmbeddingModel):
     """Vision embedding model for medical image inputs.
 
     Converts medical images to sequences of patch embeddings suitable for
@@ -115,11 +116,13 @@ def __init__(
         freeze_backbone: bool = False,
         dropout: float = 0.0,
         use_cls_token: bool = False,
+        pool: Optional[Literal["mean"]] = None,
     ) -> None:
         super().__init__(dataset)
 
-        self.embedding_dim = embedding_dim
+        self._embedding_dim = embedding_dim
         self.patch_size = patch_size
+        self.pool = pool
         self.backbone_type = backbone
         self.use_cls_token = use_cls_token
 
@@ -157,6 +160,10 @@ def __init__(
                 "in_channels": in_channels,
             }
 
+    @property
+    def embedding_dim(self) -> int:
+        return self._embedding_dim
+
     def _infer_channels(self, processor: ImageProcessor) -> int:
         """Infer number of input channels from processor mode."""
         mode = getattr(processor, "mode", None)
@@ -179,7 +186,7 @@ def _build_embedding_layer(
         if backbone == "patch":
             num_patches = (image_size // self.patch_size) ** 2
             self.embedding_layers[field_name] = PatchEmbedding(
-                image_size, self.patch_size, in_channels, self.embedding_dim
+                image_size, self.patch_size, in_channels, self._embedding_dim
             )
 
         elif backbone == "cnn":
@@ -191,8 +198,8 @@ def _build_embedding_layer(
                 nn.Conv2d(64, 128, 3, stride=2, padding=1),
                 nn.BatchNorm2d(128),
                 nn.ReLU(inplace=True),
-                nn.Conv2d(128, self.embedding_dim, 3, stride=2, padding=1),
-                nn.BatchNorm2d(self.embedding_dim),
+                nn.Conv2d(128, self._embedding_dim, 3, stride=2, padding=1),
+                nn.BatchNorm2d(self._embedding_dim),
                 nn.ReLU(inplace=True),
                 nn.AdaptiveAvgPool2d((7, 7)),
                 nn.Flatten(2),
@@ -242,7 +249,7 @@ def _build_resnet_backbone(
 
         return nn.Sequential(
             backbone_net,
-            nn.Conv2d(feature_dim, self.embedding_dim, kernel_size=1),
+            nn.Conv2d(feature_dim, self._embedding_dim, kernel_size=1),
             nn.Flatten(2),
             Permute(0, 2, 1),
         )
@@ -281,6 +288,9 @@ def forward(
             x = x + self.pos_embeddings[field_name]
             x = self.dropout(x)
 
+            if self.pool == "mean":
+                x = x.mean(dim=1, keepdim=True)
+
             embedded[field_name] = x
 
             if output_mask:
@@ -296,16 +306,19 @@ def get_output_info(self, field_name: str) -> Dict[str, Any]:
             raise KeyError(f"Field '{field_name}' not found")
 
         info = self._field_info[field_name].copy()
-        info["embedding_dim"] = self.embedding_dim
+        info["embedding_dim"] = self._embedding_dim
         info["has_cls_token"] = self.use_cls_token
-        info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0)
+        if self.pool == "mean":
+            info["num_tokens"] = 1
+        else:
+            info["num_tokens"] = info["num_patches"] + (1 if self.use_cls_token else 0)
         return info
 
     def __repr__(self) -> str:
         fields = list(self.embedding_layers.keys())
         return (
             f"VisionEmbeddingModel(backbone={self.backbone_type!r}, "
-            f"embedding_dim={self.embedding_dim}, fields={fields})"
+            f"embedding_dim={self._embedding_dim}, fields={fields})"
         )
 
 
@@ -345,14 +358,27 @@ def __repr__(self) -> str:
         use_cls_token=True,
     )
 
+    model_pooled = VisionEmbeddingModel(
+        dataset=dataset,
+        embedding_dim=128,
+        backbone="cnn",
+        pool="mean",
+    )
+
+
+
     loader = get_dataloader(dataset, batch_size=4, shuffle=False)
     batch = next(iter(loader))
 
+    embeddings_pooled = model_pooled({"chest_xray": batch["chest_xray"]})
+    print(f"Pooled output shape: {embeddings_pooled['chest_xray'].shape}")  # expect (4, 1, 128)
+    print(f"Pooled output info: {model_pooled.get_output_info('chest_xray')}")  # expect num_tokens=1
+
+
     embeddings = model({"chest_xray": batch["chest_xray"]})
     print(f"Input shape: {batch['chest_xray'].shape}")
     print(f"Output shape: {embeddings['chest_xray'].shape}")
     print(f"Output info: {model.get_output_info('chest_xray')}")
 
     # Cleanup
-   
-    shutil.rmtree(temp_dir)
\ No newline at end of file
+    shutil.rmtree(temp_dir)
diff --git a/pyhealth/models/jamba_ehr.py b/pyhealth/models/jamba_ehr.py
index fea902bd1..37d738f3f 100644
--- a/pyhealth/models/jamba_ehr.py
+++ b/pyhealth/models/jamba_ehr.py
@@ -15,6 +15,7 @@
 from pyhealth.datasets import SampleDataset
 from pyhealth.models import BaseModel
 from pyhealth.models.embedding import EmbeddingModel
+from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel
 from pyhealth.models.transformer import TransformerBlock
 from pyhealth.models.ehrmamba import MambaBlock
 from pyhealth.models.utils import get_last_visit
@@ -177,6 +178,11 @@ class JambaEHR(BaseModel):
     by an independent :class:`JambaLayer`. The resulting patient embeddings
     are concatenated and projected through a classification head.
 
+    When ``unified_embedding`` is supplied the model switches to **unified
+    mode**: all temporal fields are jointly embedded and time-sorted by
+    :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single*
+    :class:`JambaLayer` rather than one layer per field.
+
     Args:
         dataset (SampleDataset): Dataset providing processed inputs.
         embedding_dim (int): Embedding and hidden dimension. Default 128.
@@ -186,6 +192,8 @@ class JambaEHR(BaseModel):
         dropout (float): Dropout rate. Default 0.3.
         state_size (int): SSM state size in Mamba blocks. Default 16.
         conv_kernel (int): Causal conv kernel in Mamba blocks. Default 4.
+        unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when
+            provided, enables unified multi-modal mode with a single JambaLayer.
 
     Examples:
         >>> from pyhealth.datasets import create_sample_dataset, get_dataloader
@@ -234,6 +242,7 @@ def __init__(
         dropout: float = 0.3,
         state_size: int = 16,
         conv_kernel: int = 4,
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
     ):
         super(JambaEHR, self).__init__(dataset=dataset)
         self.embedding_dim = embedding_dim
@@ -243,6 +252,7 @@ def __init__(
         self.dropout_rate = dropout
         self.state_size = state_size
         self.conv_kernel = conv_kernel
+        self._use_unified = unified_embedding is not None
 
         assert (
             len(self.label_keys) == 1
@@ -250,11 +260,12 @@ def __init__(
         self.label_key = self.label_keys[0]
         self.mode = self.dataset.output_schema[self.label_key]
 
-        self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+        output_size = self.get_output_size()
+        self.dropout = nn.Dropout(dropout)
 
-        self.jamba: nn.ModuleDict = nn.ModuleDict()
-        for feature_key in self.feature_keys:
-            self.jamba[feature_key] = JambaLayer(
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            self._unified_jamba = JambaLayer(
                 feature_size=embedding_dim,
                 num_transformer_layers=num_transformer_layers,
                 num_mamba_layers=num_mamba_layers,
@@ -263,12 +274,66 @@ def __init__(
                 state_size=state_size,
                 conv_kernel=conv_kernel,
             )
+            self.fc = nn.Linear(embedding_dim, output_size)
+        else:
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            self.jamba: nn.ModuleDict = nn.ModuleDict()
+            for feature_key in self.feature_keys:
+                self.jamba[feature_key] = JambaLayer(
+                    feature_size=embedding_dim,
+                    num_transformer_layers=num_transformer_layers,
+                    num_mamba_layers=num_mamba_layers,
+                    heads=heads,
+                    dropout=dropout,
+                    state_size=state_size,
+                    conv_kernel=conv_kernel,
+                )
+            self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size)
+
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Build the inputs dict required by UnifiedMultimodalEmbeddingModel."""
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+        return inputs
+
+    def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode.
+
+        Calls UnifiedMultimodalEmbeddingModel to produce a single
+        temporally-sorted event sequence, then encodes it with one shared
+        JambaLayer and pools to the last valid event.
+        """
+        inputs = self._build_unified_inputs(kwargs)
+        out = self.embedding_model(inputs)
+        sequence = out["sequence"]          # (B, S_total, E)
+        mask     = out["mask"]              # (B, S_total) float, 1=valid 0=pad
 
-        output_size = self.get_output_size()
-        self.dropout = nn.Dropout(dropout)
-        self.fc = nn.Linear(
-            len(self.feature_keys) * embedding_dim, output_size
-        )
+        _, cls_emb = self._unified_jamba(sequence, mask)
+        logits = self.fc(self.dropout(cls_emb))
+        y_prob = self.prepare_y_prob(logits)
+
+        results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob}
+        if self.label_key in kwargs:
+            y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device)
+            results["loss"]   = self.get_loss_function()(logits, y_true)
+            results["y_true"] = y_true
+        if kwargs.get("embed", False):
+            results["embed"] = cls_emb
+        return results
 
     @staticmethod
     def _pool_embedding(x: torch.Tensor) -> torch.Tensor:
@@ -317,9 +382,10 @@ def forward(
     ) -> Dict[str, torch.Tensor]:
         """Forward propagation.
 
-        Embeds each feature stream, encodes through the hybrid
-        Transformer-Mamba stack, concatenates per-stream patient
-        representations, and projects to label space.
+        In **unified mode** (when ``unified_embedding`` was supplied at init)
+        the model jointly embeds all temporal fields and processes them with a
+        single JambaLayer backbone.  Otherwise each field is embedded and
+        encoded independently.
 
         Args:
             **kwargs: Must include all feature keys (tensors or tuples
@@ -330,6 +396,9 @@ def forward(
                 ``y_prob``, ``y_true``, ``logit``, and optionally
                 ``embed`` if ``kwargs["embed"] is True``.
         """
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
         patient_emb = []
 
         for feature_key in self.feature_keys:
diff --git a/pyhealth/models/mlp.py b/pyhealth/models/mlp.py
index 299dc151e..ab15432f0 100644
--- a/pyhealth/models/mlp.py
+++ b/pyhealth/models/mlp.py
@@ -1,4 +1,4 @@
-from typing import Dict, cast
+from typing import Any, Dict, Optional, cast
 
 import torch
 import torch.nn as nn
@@ -8,6 +8,7 @@
 from pyhealth.interpret.api import Interpretable
 
 from .embedding import EmbeddingModel
+from .embedding.unified import UnifiedMultimodalEmbeddingModel
 
 
 class MLP(BaseModel, Interpretable):
@@ -110,12 +111,14 @@ def __init__(
         hidden_dim: int = 128,
         n_layers: int = 2,
         activation: str = "relu",
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
         **kwargs,
     ):
         super(MLP, self).__init__(dataset)
         self.embedding_dim = embedding_dim
         self.hidden_dim = hidden_dim
         self.n_layers = n_layers
+        self._use_unified = unified_embedding is not None
 
         # validate kwargs for MLP layer
         if "input_size" in kwargs:
@@ -126,9 +129,6 @@ def __init__(
         assert len(self.label_keys) == 1, "Only one label key is supported"
         self.label_key = self.label_keys[0]
 
-        # Use the EmbeddingModel to handle embedding logic
-        self.embedding_model = EmbeddingModel(dataset, embedding_dim)
-
         # Set up activation function
         if activation == "relu":
             self.activation = nn.ReLU()
@@ -143,18 +143,74 @@ def __init__(
         else:
             raise ValueError(f"Unsupported activation function {activation}")
 
-        # Create MLP layers for each feature
-        self.mlp = nn.ModuleDict()
-        for feature_key in self.feature_keys:
-            Modules = []
-            Modules.append(nn.Linear(self.embedding_dim, self.hidden_dim))
-            for _ in range(self.n_layers - 1):
-                Modules.append(self.activation)
-                Modules.append(nn.Linear(self.hidden_dim, self.hidden_dim))
-            self.mlp[feature_key] = nn.Sequential(*Modules)
-
         output_size = self.get_output_size()
-        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            modules = [nn.Linear(embedding_dim, hidden_dim)]
+            for _ in range(n_layers - 1):
+                modules.extend([self.activation, nn.Linear(hidden_dim, hidden_dim)])
+            self.mlp = nn.ModuleDict({"unified": nn.Sequential(*modules)})
+            self.fc = nn.Linear(hidden_dim, output_size)
+        else:
+            # Use the EmbeddingModel to handle embedding logic
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            # Create MLP layers for each feature
+            self.mlp = nn.ModuleDict()
+            for feature_key in self.feature_keys:
+                modules = [nn.Linear(self.embedding_dim, self.hidden_dim)]
+                for _ in range(self.n_layers - 1):
+                    modules.extend([self.activation, nn.Linear(self.hidden_dim, self.hidden_dim)])
+                self.mlp[feature_key] = nn.Sequential(*modules)
+            self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Build the inputs dict required by UnifiedMultimodalEmbeddingModel."""
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+        return inputs
+
+    def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode.
+
+        Embeds all temporal fields jointly, mean-pools the event sequence,
+        applies a single MLP, and projects to label space.
+        """
+        inputs = self._build_unified_inputs(kwargs)
+        out = self.embedding_model(inputs)
+        sequence = out["sequence"]          # (B, S, E)
+        mask = out["mask"].float()          # (B, S)
+
+        # Masked mean-pool over the event sequence
+        x = (sequence * mask.unsqueeze(-1)).sum(dim=1)
+        x = x / mask.sum(dim=1, keepdim=True).clamp(min=1)  # (B, E)
+
+        x = self.mlp["unified"](x)          # (B, hidden_dim)
+        logits = self.fc(x)
+        y_prob = self.prepare_y_prob(logits)
+
+        results: Dict[str, torch.Tensor] = {"logit": logits, "y_prob": y_prob}
+        if self.label_key in kwargs:
+            y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device)
+            results["loss"] = self.get_loss_function()(logits, y_true)
+            results["y_true"] = y_true
+        if kwargs.get("embed", False):
+            results["embed"] = x
+        return results
 
     @staticmethod
     def mean_pooling(x, mask):
@@ -309,6 +365,11 @@ def forward(
     ) -> Dict[str, torch.Tensor]:
         """Forward propagation.
 
+        In **unified mode** (when ``unified_embedding`` was supplied at init)
+        the model jointly embeds all temporal fields, mean-pools the event
+        sequence, and processes it with a single MLP.  Otherwise each field
+        is embedded and encoded independently.
+
         Args:
             **kwargs: keyword arguments for the model.
 
@@ -326,6 +387,9 @@ def forward(
                 logit: the raw logits before activation.
                 embed: (if embed=True in kwargs) the patient embedding.
         """
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
         for feature_key in self.feature_keys:
             feature = kwargs[feature_key]
 
@@ -355,10 +419,9 @@ def forward(
                 batch_size, seq_len, inner_len = value.shape
                 value = value.view(batch_size, seq_len * inner_len)
                 if mask is not None:
-                     mask = mask.to(self.device)
-                     # Flatten mask properly if it exists
-                     if mask.dim() == 3:
-                         mask = mask.view(batch_size, seq_len * inner_len)
+                    mask = mask.to(self.device)
+                    if mask.dim() == 3:
+                        mask = mask.view(batch_size, seq_len * inner_len)
 
             if mask is not None:
                 mask = mask.to(self.device)
diff --git a/pyhealth/models/rnn.py b/pyhealth/models/rnn.py
index 4c9ba0550..b131628a9 100644
--- a/pyhealth/models/rnn.py
+++ b/pyhealth/models/rnn.py
@@ -1,4 +1,4 @@
-from typing import Dict, Optional, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import torch
 import torch.nn as nn
@@ -20,6 +20,7 @@
 )
 
 from .embedding import EmbeddingModel
+from .embedding.unified import UnifiedMultimodalEmbeddingModel
 
 
 class RNNLayer(nn.Module):
@@ -206,6 +207,7 @@ def __init__(
         dataset: SampleDataset,
         embedding_dim: int = 128,
         hidden_dim: int = 128,
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
         **kwargs
     ):
         super(RNN, self).__init__(
@@ -213,6 +215,7 @@ def __init__(
         )
         self.embedding_dim = embedding_dim
         self.hidden_dim = hidden_dim
+        self._use_unified = unified_embedding is not None
         # validate kwargs for RNN layer
         if "input_size" in kwargs:
             raise ValueError("input_size is determined by embedding_dim")
@@ -222,20 +225,71 @@ def __init__(
         self.label_key = self.label_keys[0]
         self.mode = self.dataset.output_schema[self.label_key]
 
-        self.embedding_model = EmbeddingModel(dataset, embedding_dim)
-
-        self.rnn = nn.ModuleDict()
-        for feature_key in self.dataset.input_processors.keys():
-            self.rnn[feature_key] = RNNLayer(
-                input_size=embedding_dim, hidden_size=hidden_dim, **kwargs
-            )
         output_size = self.get_output_size()
-        self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            self.rnn = nn.ModuleDict({
+                "unified": RNNLayer(input_size=embedding_dim, hidden_size=hidden_dim, **kwargs)
+            })
+            self.fc = nn.Linear(hidden_dim, output_size)
+        else:
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            self.rnn = nn.ModuleDict()
+            for feature_key in self.dataset.input_processors.keys():
+                self.rnn[feature_key] = RNNLayer(
+                    input_size=embedding_dim, hidden_size=hidden_dim, **kwargs
+                )
+            self.fc = nn.Linear(len(self.feature_keys) * self.hidden_dim, output_size)
+
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Build the inputs dict required by UnifiedMultimodalEmbeddingModel."""
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+        return inputs
+
+    def _forward_unified(self, **kwargs: Any) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode.
+
+        Embeds all temporal fields jointly as a single time-sorted sequence
+        and processes it with one RNN backbone.
+        """
+        inputs = self._build_unified_inputs(kwargs)
+        out = self.embedding_model(inputs)
+        sequence = out["sequence"]          # (B, S, E)
+        mask = out["mask"].int()            # (B, S)
+
+        _, last_hidden = self.rnn["unified"](sequence, mask)  # (B, hidden_dim)
+        logits = self.fc(last_hidden)
+        y_true = kwargs[self.label_key].to(self.device)
+        loss = self.get_loss_function()(logits, y_true)
+        y_prob = self.prepare_y_prob(logits)
+        results = {"loss": loss, "y_prob": y_prob, "y_true": y_true, "logit": logits}
+        if kwargs.get("embed", False):
+            results["embed"] = last_hidden
+        return results
 
     def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
         """Forward propagation.
 
-        The label `kwargs[self.label_key]` is a list of labels for each patient.
+        In **unified mode** (when ``unified_embedding`` was supplied at init)
+        the model jointly embeds all temporal fields as a single time-sorted
+        sequence and processes it with one RNN.  Otherwise each field is
+        embedded and encoded independently.
 
         Args:
             **kwargs: keyword arguments for the model. The keys must contain
@@ -249,6 +303,9 @@ def forward(self, **kwargs) -> Dict[str, torch.Tensor]:
                 - logit: a tensor representing the logits.
                 - embed (optional): a tensor representing the patient embeddings if requested.
         """
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
         patient_emb = []
         
         # We need to preprocess kwargs to extract values and masks for EmbeddingModel
diff --git a/pyhealth/models/transformer.py b/pyhealth/models/transformer.py
index cc0dfc5ca..f678403b2 100644
--- a/pyhealth/models/transformer.py
+++ b/pyhealth/models/transformer.py
@@ -12,6 +12,7 @@
 from pyhealth.datasets import SampleDataset
 from pyhealth.models import BaseModel
 from pyhealth.models.embedding import EmbeddingModel
+from pyhealth.models.embedding.unified import UnifiedMultimodalEmbeddingModel
 from pyhealth.interpret.api import CheferInterpretable
 
 # VALID_OPERATION_LEVEL = ["visit", "event"]
@@ -54,7 +55,7 @@ def forward(
             # Use -inf so softmax produces exact zeros on padded positions,
             # avoiding a second masked_fill after softmax (saves one full
             # [B, H, S, S] boolean allocation and an extra copy).
-            pad_mask = (mask == 0)
+            pad_mask = mask == 0
             scores = scores.masked_fill(pad_mask, -1e9)
         p_attn = self.softmax(scores)
         if dropout is not None:
@@ -164,7 +165,7 @@ def forward(
             self.attn_map = None
         # 3) "Concat" using a view and apply a final linear.
         x = x.transpose(1, 2).contiguous().view(batch_size, -1, self.h * self.d_k)
-  
+
         return self.output_linear(x)
 
 
@@ -246,7 +247,7 @@ def set_activation_hooks(self, hooks) -> None:
         """Deprecated compatibility stub; no-op."""
         return None
 
-    def forward(self, x, mask=None, register_hook = False):
+    def forward(self, x, mask=None, register_hook=False):
         """Forward propagation.
 
         Args:
@@ -256,7 +257,12 @@ def forward(self, x, mask=None, register_hook = False):
         Returns:
             A tensor of shape [batch_size, seq_len, hidden]
         """
-        x = self.input_sublayer(x, lambda _x: self.attention(_x, _x, _x, mask=mask, register_hook=register_hook))
+        x = self.input_sublayer(
+            x,
+            lambda _x: self.attention(
+                _x, _x, _x, mask=mask, register_hook=register_hook
+            ),
+        )
         x = self.output_sublayer(x, lambda _x: self.feed_forward(_x, mask=mask))
         return self.dropout(x)
 
@@ -297,7 +303,10 @@ def set_activation_hooks(self, hooks) -> None:
         return None
 
     def forward(
-        self, x: torch.Tensor, mask: Optional[torch.Tensor] = None, register_hook: bool = False
+        self,
+        x: torch.Tensor,
+        mask: Optional[torch.Tensor] = None,
+        register_hook: bool = False,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Forward propagation.
 
@@ -328,12 +337,21 @@ class Transformer(BaseModel, CheferInterpretable):
     an independent :class:`TransformerLayer`. The resulting [CLS]-style
     embeddings are concatenated and passed to a classification head.
 
+    When ``unified_embedding`` is supplied the model switches to **unified
+    mode**: all temporal fields are jointly embedded and time-sorted by
+    :class:`UnifiedMultimodalEmbeddingModel`, then processed by a *single*
+    :class:`TransformerLayer` rather than one layer per field.  This allows
+    full cross-modal attention over the interleaved event sequence.
+
     Args:
         dataset (SampleDataset): dataset providing processed inputs.
         embedding_dim (int): shared embedding dimension.
         heads (int): number of attention heads per transformer block.
         dropout (float): dropout rate applied inside transformer blocks.
         num_layers (int): number of transformer blocks per feature stream.
+        unified_embedding (UnifiedMultimodalEmbeddingModel, optional): when
+            provided, the model uses a single backbone over the unified
+            multi-modal sequence instead of per-field transformers.
 
     Examples:
         >>> from pyhealth.datasets import create_sample_dataset, get_dataloader
@@ -377,6 +395,7 @@ def __init__(
         dropout: float = 0.5,
         num_layers: int = 1,
         max_seq_len: int = 1024,
+        unified_embedding: Optional[UnifiedMultimodalEmbeddingModel] = None,
     ):
         super().__init__(dataset=dataset)
         self.embedding_dim = embedding_dim
@@ -385,6 +404,7 @@ def __init__(
         self.num_layers = num_layers
         self.max_seq_len = max_seq_len
         self._attention_hooks_enabled = False
+        self._use_unified = unified_embedding is not None
 
         assert (
             len(self.label_keys) == 1
@@ -392,19 +412,28 @@ def __init__(
         self.label_key = self.label_keys[0]
         self.mode = self.dataset.output_schema[self.label_key]
 
-        self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+        output_size = self.get_output_size()
 
-        self.transformer: nn.ModuleDict = nn.ModuleDict()
-        for feature_key in self.feature_keys:
-            self.transformer[feature_key] = TransformerLayer(
+        if self._use_unified:
+            self.embedding_model = unified_embedding
+            self._unified_backbone = TransformerLayer(
                 feature_size=embedding_dim,
                 heads=heads,
                 dropout=dropout,
                 num_layers=num_layers,
             )
-
-        output_size = self.get_output_size()
-        self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size)
+            self.fc = nn.Linear(embedding_dim, output_size)
+        else:
+            self.embedding_model = EmbeddingModel(dataset, embedding_dim)
+            self.transformer: nn.ModuleDict = nn.ModuleDict()
+            for feature_key in self.feature_keys:
+                self.transformer[feature_key] = TransformerLayer(
+                    feature_size=embedding_dim,
+                    heads=heads,
+                    dropout=dropout,
+                    num_layers=num_layers,
+                )
+            self.fc = nn.Linear(len(self.feature_keys) * embedding_dim, output_size)
 
     def _pool_embedding(self, x: torch.Tensor) -> torch.Tensor:
         """Pool nested embeddings to ``[batch, seq_len, hidden]`` format.
@@ -443,6 +472,61 @@ def _mask_from_embeddings(x: torch.Tensor) -> torch.Tensor:
             mask[invalid_rows, 0] = True
         return mask.bool()
 
+    def _build_unified_inputs(
+        self, kwargs: Dict[str, Any]
+    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        """Build inputs expected by UnifiedMultimodalEmbeddingModel."""
+
+        inputs: Dict[str, Dict[str, torch.Tensor]] = {}
+        for field_name in self.feature_keys:
+            feature = kwargs[field_name]
+            if isinstance(feature, torch.Tensor):
+                feature = (feature,)
+
+            schema = self.dataset.input_processors[field_name].schema()
+            field_dict: Dict[str, torch.Tensor] = {}
+            if "value" in schema:
+                field_dict["value"] = feature[schema.index("value")].to(self.device)
+            if "time" in schema:
+                field_dict["time"] = feature[schema.index("time")].to(self.device)
+            if "mask" in schema:
+                field_dict["mask"] = feature[schema.index("mask")].to(self.device)
+            inputs[field_name] = field_dict
+
+        return inputs
+
+    def _forward_unified(
+        self,
+        **kwargs: torch.Tensor | tuple[torch.Tensor, ...],
+    ) -> Dict[str, torch.Tensor]:
+        """Forward pass in unified-embedding mode."""
+
+        register_hook = self._attention_hooks_enabled
+        inputs = self._build_unified_inputs(cast(Dict[str, Any], kwargs))
+        out = self.embedding_model(inputs)
+        sequence = cast(torch.Tensor, out["sequence"])
+        event_mask = cast(torch.Tensor, out["mask"]).bool()
+
+        _, patient_emb = self._unified_backbone(sequence, event_mask, register_hook)
+
+        logits = self.fc(patient_emb)
+        y_prob = self.prepare_y_prob(logits)
+
+        results: Dict[str, torch.Tensor] = {
+            "logit": logits,
+            "y_prob": y_prob,
+        }
+
+        if self.label_key in kwargs:
+            y_true = cast(torch.Tensor, kwargs[self.label_key]).to(self.device)
+            loss = self.get_loss_function()(logits, y_true)
+            results["loss"] = loss
+            results["y_true"] = y_true
+
+        if kwargs.get("embed", False):
+            results["embed"] = patient_emb
+        return results
+
     def forward_from_embedding(
         self,
         **kwargs: torch.Tensor | tuple[torch.Tensor, ...],
@@ -500,8 +584,7 @@ def forward_from_embedding(
 
             if value is None:
                 raise ValueError(
-                    f"Feature '{feature_key}' must contain 'value' "
-                    f"in the schema."
+                    f"Feature '{feature_key}' must contain 'value' " f"in the schema."
                 )
             else:
                 value = value.to(self.device)
@@ -515,9 +598,7 @@ def forward_from_embedding(
             else:
                 mask = self._mask_from_embeddings(value).to(self.device)
 
-            _, cls_emb = self.transformer[feature_key](
-                value, mask, register_hook
-            )
+            _, cls_emb = self.transformer[feature_key](value, mask, register_hook)
             patient_emb.append(cls_emb)
 
         patient_emb = torch.cat(patient_emb, dim=1)
@@ -545,6 +626,11 @@ def forward(
     ) -> Dict[str, torch.Tensor]:
         """Forward propagation.
 
+        In **unified mode** (when ``unified_embedding`` was supplied at init)
+        the model jointly embeds all temporal fields and processes them with a
+        single transformer backbone.  Otherwise each field is embedded and
+        encoded independently.
+
         Args:
             **kwargs: keyword arguments for the model.
 
@@ -562,6 +648,9 @@ def forward(
                 logit: the raw logits before activation.
                 embed: (if embed=True in kwargs) the patient embedding.
         """
+        if self._use_unified:
+            return self._forward_unified(**kwargs)
+
         for feature_key in self.feature_keys:
             feature = kwargs[feature_key]
 
@@ -575,15 +664,16 @@ def forward(
 
             if value is None:
                 raise ValueError(
-                    f"Feature '{feature_key}' must contain 'value' "
-                    f"in the schema."
+                    f"Feature '{feature_key}' must contain 'value' " f"in the schema."
                 )
             else:
                 value = value.to(self.device)
 
             if mask is not None:
                 mask = mask.to(self.device)
-                value = self.embedding_model({feature_key: value}, masks={feature_key: mask})[feature_key]
+                value = self.embedding_model(
+                    {feature_key: value}, masks={feature_key: mask}
+                )[feature_key]
             else:
                 value = self.embedding_model({feature_key: value})[feature_key]
 
@@ -591,9 +681,9 @@ def forward(
             # Reconstruct tuple with embedded value
             # Note: we need to handle list/tuple conversion carefully
             # feature is a tuple.
-            
+
             # Simple slice reconstruction
-            kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1:]
+            kwargs[feature_key] = feature[:i] + (value,) + feature[i + 1 :]
 
         return self.forward_from_embedding(**kwargs)
 
@@ -621,9 +711,7 @@ def get_attention_layers(
                     cast(TransformerBlock, blk).attention.get_attn_map(),
                     cast(TransformerBlock, blk).attention.get_attn_grad(),
                 )
-                for blk in cast(
-                    TransformerLayer, self.transformer[key]
-                ).transformer
+                for blk in cast(TransformerLayer, self.transformer[key]).transformer
             ]
             for key in self.feature_keys
         }
diff --git a/pyhealth/models/unified_embedding.py b/pyhealth/models/unified_embedding.py
deleted file mode 100644
index 014326b41..000000000
--- a/pyhealth/models/unified_embedding.py
+++ /dev/null
@@ -1,327 +0,0 @@
-"""UnifiedMultimodalEmbeddingModel — temporally aligned multimodal embedding.
-
-Takes K temporal features ( dict outputs from ``TemporalFeatureProcessor``
-subclasses ), embeds each event with a modality-specific encoder, then
-interleaves all events on a shared timeline by sorting on timestamp and adding
-sinusoidal time embeddings + learned modality-type embeddings.
-
-Output shape: ``(B, S_total, E')`` — a single sequence of events usable by
-any downstream sequence model (Transformer, Mamba, RNN, …).
-
-Quickstart::
-
-    from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel
-    from pyhealth.datasets.collate import collate_temporal
-    model = UnifiedMultimodalEmbeddingModel(dataset, embedding_dim=128)
-    # inside forward:
-    #   inputs = {field: {"value": Tensor, "time": Tensor, ...}, ...}
-    out = model(inputs)
-    # out["sequence"]: (B, S_total, 128)
-    # out["mask"]:     (B, S_total)      — 1 = real event, 0 = padding
-    # out["time"]:     (B, S_total)      — hours from first event
-"""
-from __future__ import annotations
-
-import math
-from typing import Any
-
-import torch
-import torch.nn as nn
-
-from pyhealth.processors.base_processor import ModalityType, TemporalFeatureProcessor
-
-
-# ── Helpers ───────────────────────────────────────────────────────────────────
-
-
-class SinusoidalTimeEmbedding(nn.Module):
-    """Continuous sinusoidal embedding for scalar time values (in hours).
-
-    Identical in spirit to the positional encoding in "Attention is All You
-    Need" but operating on real-valued timestamps rather than integer positions.
-
-    Args:
-        dim: Output embedding dimension (must be even).
-        max_hours: Maximum expected time value in hours.  Values are normalised
-            to ``[0, 2π]`` before the sin/cos projection.  Default 720 (30 days).
-
-    Shape:
-        Input:  ``(*, )``  float tensor of times in hours
-        Output: ``(*, dim)``
-    """
-
-    def __init__(self, dim: int, max_hours: float = 720.0):
-        super().__init__()
-        assert dim % 2 == 0, f"dim must be even, got {dim}"
-        self.dim = dim
-        self.max_hours = max_hours
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(10000.0) * torch.arange(half, dtype=torch.float32) / (half - 1)
-        )
-        self.register_buffer("freqs", freqs)  # (dim//2,)
-
-    def forward(self, t: torch.Tensor) -> torch.Tensor:
-        """:param t: ``(...,)`` float, times in hours."""
-        t_norm = t / self.max_hours * 2 * math.pi           # (...,)
-        args = t_norm.unsqueeze(-1) * self.freqs             # (..., dim//2)
-        return torch.cat([args.sin(), args.cos()], dim=-1)   # (..., dim)
-
-
-def _build_image_encoder(embedding_dim: int) -> nn.Module:
-    """Lightweight 5-layer CNN encoder: C × H × W → embedding_dim.
-
-    Uses ``torchvision.models.resnet18`` pre-trained backbone, strips the
-    final FC layer, and adds a projection to ``embedding_dim``.  Falls back to
-    a toy Conv-pool-flatten network if torchvision is not installed.
-    """
-    try:
-        import torchvision.models as tv
-
-        backbone = tv.resnet18(weights=None)
-        in_features = backbone.fc.in_features
-        backbone.fc = nn.Linear(in_features, embedding_dim)
-        return backbone
-    except ImportError:
-        # Minimal fallback: single conv → global avg pool → linear
-        return nn.Sequential(
-            nn.Conv2d(3, 32, 3, padding=1),
-            nn.ReLU(),
-            nn.AdaptiveAvgPool2d(1),
-            nn.Flatten(),
-            nn.Linear(32, embedding_dim),
-        )
-
-
-# ── Main model ───────────────────────────────────────────────────────────────
-
-
-class UnifiedMultimodalEmbeddingModel(nn.Module):
-    """Embed heterogeneous temporal features into a single aligned sequence.
-
-    **All** input processors must be ``TemporalFeatureProcessor`` subclasses.
-    Non-temporal processors (e.g. ``SequenceProcessor``, ``MultiHotProcessor``)
-    are rejected with a clear error — use the existing ``EmbeddingModel`` for
-    those fields.
-
-    Algorithm
-    ---------
-    For each temporal field:
-
-    1. Route ``inputs[field]["value"]`` through a modality-specific encoder →
-       ``(B, N_i, E')`` per-event embeddings.
-    2. Retrieve ``inputs[field]["time"]`` → ``(B, N_i)`` timestamps (hours).
-    3. (Optional) Retrieve ``inputs[field]["mask"]`` → ``(B, N_i, L)`` or
-       ``(B, N_i)`` attention mask; reduced to event-level ``(B, N_i)`` if
-       token-level.
-
-    Then:
-
-    4. Concatenate across all fields → ``(B, S_total, E')``.
-    5. Sort events along dim=1 by timestamp (ascending).
-    6. Add ``SinusoidalTimeEmbedding(time)`` + ``type_embedding(modality_idx)``.
-    7. Return ``{"sequence", "time", "mask", "type_ids"}``.
-
-    Args:
-        processors: ``dict[field_name, TemporalFeatureProcessor]`` — the
-            processors for each temporal field in the dataset.  Pass
-            ``dataset.input_processors`` directly.
-        embedding_dim: Shared embedding dimension ``E'``.
-        time_embedding: ``"sinusoidal"`` (default) or ``"learned"``.
-        max_time_hours: Normalisation constant for the time embedding.
-            Defaults to 720 h (30 days).
-
-    Example::
-
-        model = UnifiedMultimodalEmbeddingModel(
-            processors=dataset.input_processors,
-            embedding_dim=128,
-        )
-        # inputs: {field: {"value": Tensor, "time": Tensor, "mask": Tensor}}
-        out = model(inputs)
-        seq = out["sequence"]   # (B, S_total, 128)
-        mask = out["mask"]      # (B, S_total)  float, 1=valid 0=pad
-    """
-
-    def __init__(
-        self,
-        processors: dict[str, Any],
-        embedding_dim: int = 128,
-        time_embedding: str = "sinusoidal",
-        max_time_hours: float = 720.0,
-    ):
-        super().__init__()
-        self.embedding_dim = embedding_dim
-
-        self.encoders: nn.ModuleDict = nn.ModuleDict()
-        self.projections: nn.ModuleDict = nn.ModuleDict()
-        self.modality_types: dict[str, ModalityType] = {}
-
-        for field_name, processor in processors.items():
-            if not isinstance(processor, TemporalFeatureProcessor):
-                raise TypeError(
-                    f"UnifiedMultimodalEmbeddingModel requires every input processor "
-                    f"to be a TemporalFeatureProcessor subclass, but '{field_name}' "
-                    f"uses {type(processor).__name__}.  For non-temporal fields use "
-                    f"the existing EmbeddingModel."
-                )
-
-            m = processor.modality()
-            self.modality_types[field_name] = m
-
-            if m == ModalityType.CODE:
-                vocab_size = processor.value_dim()
-                self.encoders[field_name] = nn.Embedding(
-                    vocab_size, embedding_dim, padding_idx=0
-                )
-
-            elif m == ModalityType.TEXT:
-                if processor.is_token():
-                    from transformers import AutoModel
-
-                    bert = AutoModel.from_pretrained(processor.tokenizer_model)
-                    self.encoders[field_name] = bert
-                    hidden = bert.config.hidden_size
-                    if hidden != embedding_dim:
-                        self.projections[field_name] = nn.Linear(hidden, embedding_dim)
-                else:
-                    raise ValueError(
-                        f"TEXT processor '{field_name}' must use a tokenizer "
-                        f"(set tokenizer_model=...) to be used with "
-                        f"UnifiedMultimodalEmbeddingModel."
-                    )
-
-            elif m == ModalityType.IMAGE:
-                self.encoders[field_name] = _build_image_encoder(embedding_dim)
-
-            elif m in (ModalityType.NUMERIC, ModalityType.SIGNAL):
-                in_features = processor.value_dim()
-                self.encoders[field_name] = nn.Linear(in_features, embedding_dim)
-
-            else:
-                raise NotImplementedError(
-                    f"No encoder implemented for modality {m!r} (field '{field_name}')."
-                )
-
-        # Shared type embedding — one vector per unique modality in this dataset
-        unique_modalities = sorted(set(self.modality_types.values()))
-        self._modality_to_idx: dict[ModalityType, int] = {
-            mod: i for i, mod in enumerate(unique_modalities)
-        }
-        self.type_embedding = nn.Embedding(len(unique_modalities), embedding_dim)
-
-        # Time embedding
-        if time_embedding == "sinusoidal":
-            self.time_embed = SinusoidalTimeEmbedding(embedding_dim, max_time_hours)
-        else:
-            raise NotImplementedError("Only 'sinusoidal' time embedding is implemented.")
-
-    # ── Forward ───────────────────────────────────────────────────────────────
-
-    def forward(
-        self,
-        inputs: dict[str, dict[str, torch.Tensor]],
-    ) -> dict[str, torch.Tensor]:
-        """Encode and temporally align all temporal features.
-
-        Args:
-            inputs: ``{field_name: {"value": Tensor, "time": Tensor,
-                        "mask": Tensor (optional)}}``
-                — one dict per temporal feature, exactly as produced by
-                  ``collate_temporal``.
-
-        Returns:
-            A dict with keys:
-
-            * ``"sequence"`` — ``(B, S_total, E')``  temporally-sorted events
-            * ``"time"``     — ``(B, S_total)``       timestamps (hours)
-            * ``"mask"``     — ``(B, S_total)``       1=real event, 0=padding
-            * ``"type_ids"`` — ``(B, S_total)``       modality index per event
-        """
-        all_embeddings: list[torch.Tensor] = []
-        all_times:      list[torch.Tensor] = []
-        all_masks:      list[torch.Tensor] = []
-        all_types:      list[torch.Tensor] = []
-
-        for field_name, feat_dict in inputs.items():
-            value = feat_dict["value"]   # (B, N_i, ...) or (B, S, F)
-            time  = feat_dict["time"]    # (B, N_i)
-            mask  = feat_dict.get("mask")
-
-            if time is None:
-                # Fallback: treat every event as occurring at t=0
-                time = torch.zeros(value.shape[:2], device=value.device)
-
-            modality = self.modality_types[field_name]
-            encoder  = self.encoders[field_name]
-
-            # ── Encode ────────────────────────────────────────────────────
-            if modality == ModalityType.CODE:
-                emb = encoder(value)                           # (B, S, E')
-
-            elif modality == ModalityType.TEXT:
-                b, n, l = value.shape
-                flat_ids  = value.view(b * n, l)
-                flat_mask = mask.view(b * n, l) if mask is not None else None
-                out       = encoder(input_ids=flat_ids, attention_mask=flat_mask)
-                cls_emb   = out.last_hidden_state[:, 0, :]    # (B*N, H)
-                if field_name in self.projections:
-                    cls_emb = self.projections[field_name](cls_emb)
-                emb = cls_emb.view(b, n, -1)                  # (B, N, E')
-
-            elif modality == ModalityType.IMAGE:
-                b, n, c, h, w = value.shape
-                flat_imgs = value.view(b * n, c, h, w)
-                img_emb   = encoder(flat_imgs)                 # (B*N, E')
-                emb       = img_emb.view(b, n, -1)
-
-            else:  # NUMERIC / SIGNAL
-                emb = encoder(value)                           # (B, T, E')
-
-            # ── Build event-level validity mask ───────────────────────────
-            if mask is None:
-                event_mask = torch.ones(emb.shape[:2], device=emb.device)
-            else:
-                if mask.dim() > time.dim():
-                    # token-level (B, N, L) → event-level (B, N)
-                    event_mask = (mask.sum(dim=-1) > 0).float()
-                else:
-                    event_mask = mask.float()
-
-            # ── Modality type indices ─────────────────────────────────────
-            type_idx  = self._modality_to_idx[modality]
-            type_ids  = torch.full(
-                emb.shape[:2], type_idx, dtype=torch.long, device=emb.device
-            )
-
-            all_embeddings.append(emb)
-            all_times.append(time)
-            all_masks.append(event_mask)
-            all_types.append(type_ids)
-
-        # ── Concatenate across all fields ─────────────────────────────────
-        cat_emb   = torch.cat(all_embeddings, dim=1)   # (B, S_total, E')
-        cat_time  = torch.cat(all_times,      dim=1)   # (B, S_total)
-        cat_mask  = torch.cat(all_masks,      dim=1)   # (B, S_total)
-        cat_types = torch.cat(all_types,      dim=1)   # (B, S_total)
-
-        # ── Sort by time ──────────────────────────────────────────────────
-        sort_idx   = cat_time.argsort(dim=1)
-        cat_emb    = cat_emb.gather(
-            1, sort_idx.unsqueeze(-1).expand_as(cat_emb)
-        )
-        cat_time   = cat_time.gather(1, sort_idx)
-        cat_mask   = cat_mask.gather(1, sort_idx)
-        cat_types  = cat_types.gather(1, sort_idx)
-
-        # ── Add time + type embeddings ────────────────────────────────────
-        time_emb = self.time_embed(cat_time)             # (B, S_total, E')
-        type_emb = self.type_embedding(cat_types)        # (B, S_total, E')
-        final    = cat_emb + time_emb + type_emb         # (B, S_total, E')
-
-        return {
-            "sequence": final,      # (B, S_total, E')
-            "time":     cat_time,   # (B, S_total)
-            "mask":     cat_mask,   # (B, S_total)
-            "type_ids": cat_types,  # (B, S_total)
-        }
diff --git a/pyhealth/processors/label_processor.py b/pyhealth/processors/label_processor.py
index 969721995..b4080f8de 100644
--- a/pyhealth/processors/label_processor.py
+++ b/pyhealth/processors/label_processor.py
@@ -1,4 +1,5 @@
 import logging
+import os
 from typing import Any, Dict, Iterable
 
 import torch
@@ -21,12 +22,25 @@ def __init__(self):
 
     def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
         all_labels = set([sample[field] for sample in samples])
+        allow_single_class = (
+            os.environ.get("PYHEALTH_ALLOW_SINGLE_CLASS_BINARY", "0") == "1"
+        )
         if len(all_labels) != 2:
-            raise ValueError(f"Expected 2 unique labels, got {len(all_labels)}")
+            if not (allow_single_class and len(all_labels) == 1):
+                raise ValueError(f"Expected 2 unique labels, got {len(all_labels)}")
+            logger.warning(
+                "BinaryLabelProcessor received single-class labels for '%s'. "
+                "Proceeding due to PYHEALTH_ALLOW_SINGLE_CLASS_BINARY=1.",
+                field,
+            )
         if all_labels == {0, 1}:
             self.label_vocab = {0: 0, 1: 1}
         elif all_labels == {False, True}:
             self.label_vocab = {False: 0, True: 1}
+        elif all_labels == {0} or all_labels == {False}:
+            self.label_vocab = {list(all_labels)[0]: 0}
+        elif all_labels == {1} or all_labels == {True}:
+            self.label_vocab = {list(all_labels)[0]: 1}
         else:
             all_labels = list(all_labels)
             all_labels.sort()
diff --git a/pyhealth/processors/stagenet_processor.py b/pyhealth/processors/stagenet_processor.py
index 604376ec1..f9b242f0d 100644
--- a/pyhealth/processors/stagenet_processor.py
+++ b/pyhealth/processors/stagenet_processor.py
@@ -165,14 +165,29 @@ def process(
             # Flat codes: ["code1", "code2"]
             value_tensor = self._encode_codes(value_data)
 
-        # Process time if present
-        time_tensor = None
-        if time_data is not None and len(time_data) > 0:
+        # Ensure temporal output is always a tensor for serialization stability.
+        # Length must align with the first axis of value_tensor.
+        seq_len = int(value_tensor.shape[0]) if value_tensor.ndim > 0 else 1
+        if time_data is None or len(time_data) == 0:
+            time_tensor = torch.zeros(seq_len, dtype=torch.float)
+        else:
             # Handle both [0.0, 1.5] and [[0.0], [1.5]] formats
             if isinstance(time_data[0], list):
-                # Flatten [[0.0], [1.5]] -> [0.0, 1.5]
-                time_data = [t[0] if isinstance(t, list) else t for t in time_data]
-            time_tensor = torch.tensor(time_data, dtype=torch.float)
+                time_data = [t[0] if isinstance(t, list) and len(t) > 0 else t for t in time_data]
+
+            cleaned_times: List[float] = []
+            for t in time_data:
+                try:
+                    cleaned_times.append(float(t))
+                except (TypeError, ValueError):
+                    cleaned_times.append(0.0)
+
+            if len(cleaned_times) < seq_len:
+                cleaned_times.extend([0.0] * (seq_len - len(cleaned_times)))
+            elif len(cleaned_times) > seq_len:
+                cleaned_times = cleaned_times[:seq_len]
+
+            time_tensor = torch.tensor(cleaned_times, dtype=torch.float)
 
         return (time_tensor, value_tensor)
 
@@ -424,14 +439,29 @@ def process(
         # Convert to float tensor
         value_tensor = torch.tensor(value_array, dtype=torch.float)
 
-        # Process time if present
-        time_tensor = None
-        if time_data is not None and len(time_data) > 0:
+        # Ensure temporal output is always a tensor for serialization stability.
+        # Length must align with the first axis of value_tensor.
+        seq_len = int(value_tensor.shape[0]) if value_tensor.ndim > 0 else 1
+        if time_data is None or len(time_data) == 0:
+            time_tensor = torch.zeros(seq_len, dtype=torch.float)
+        else:
             # Handle both [0.0, 1.5] and [[0.0], [1.5]] formats
             if isinstance(time_data[0], list):
-                # Flatten [[0.0], [1.5]] -> [0.0, 1.5]
-                time_data = [t[0] if isinstance(t, list) else t for t in time_data]
-            time_tensor = torch.tensor(time_data, dtype=torch.float)
+                time_data = [t[0] if isinstance(t, list) and len(t) > 0 else t for t in time_data]
+
+            cleaned_times: List[float] = []
+            for t in time_data:
+                try:
+                    cleaned_times.append(float(t))
+                except (TypeError, ValueError):
+                    cleaned_times.append(0.0)
+
+            if len(cleaned_times) < seq_len:
+                cleaned_times.extend([0.0] * (seq_len - len(cleaned_times)))
+            elif len(cleaned_times) > seq_len:
+                cleaned_times = cleaned_times[:seq_len]
+
+            time_tensor = torch.tensor(cleaned_times, dtype=torch.float)
 
         return (time_tensor, value_tensor)
 
@@ -499,4 +529,4 @@ def __repr__(self):
         return (
             f"StageNetTensorProcessor(is_nested={self._is_nested}, "
             f"feature_dim={self._size})"
-        )
+        )
\ No newline at end of file
diff --git a/pyhealth/processors/time_image_processor.py b/pyhealth/processors/time_image_processor.py
index 9d313e6bc..421998d07 100644
--- a/pyhealth/processors/time_image_processor.py
+++ b/pyhealth/processors/time_image_processor.py
@@ -75,6 +75,11 @@ class TimeImageProcessor(TemporalFeatureProcessor):
             patient has more images, the most recent (by timestamp)
             are kept. If None, all images are kept. Defaults to
             None.
+        padding: Sentinel string that marks a missing image. When
+            a path equals this value, a zero tensor of shape
+            (C, H, W) is returned instead of loading from disk.
+            If None, all paths are treated as real file paths.
+            Defaults to None.
 
     Raises:
         ValueError: If normalize is True but mean or std is missing.
@@ -107,6 +112,7 @@ def __init__(
         std: Optional[List[float]] = None,
         mode: Optional[str] = None,
         max_images: Optional[int] = None,
+        padding: Optional[str] = None,
     ) -> None:
         self.image_size = image_size
         self.to_tensor = to_tensor
@@ -115,18 +121,14 @@ def __init__(
         self.std = std
         self.mode = mode
         self.max_images = max_images
+        self.padding = padding
         self.n_channels = None
 
-        if self.normalize and (
-            self.mean is None or self.std is None
-        ):
+        if self.normalize and (self.mean is None or self.std is None):
             raise ValueError(
-                "Normalization requires both mean and std to be "
-                "provided."
+                "Normalization requires both mean and std to be " "provided."
             )
-        if not self.normalize and (
-            self.mean is not None or self.std is not None
-        ):
+        if not self.normalize and (self.mean is not None or self.std is not None):
             raise ValueError(
                 "Mean and std are provided but normalize is set "
                 "to False. Either provide normalize=True, or "
@@ -146,36 +148,49 @@ def _build_transform(self) -> transforms.Compose:
         transform_list = []
         if self.mode is not None:
             transform_list.append(
-                transforms.Lambda(
-                    partial(_convert_mode, mode=self.mode)
-                )
+                transforms.Lambda(partial(_convert_mode, mode=self.mode))
             )
         if self.image_size is not None:
-            transform_list.append(
-                transforms.Resize(
-                    (self.image_size, self.image_size)
-                )
-            )
+            transform_list.append(transforms.Resize((self.image_size, self.image_size)))
         if self.to_tensor:
             transform_list.append(transforms.ToTensor())
         if self.normalize:
-            transform_list.append(
-                transforms.Normalize(
-                    mean=self.mean, std=self.std
-                )
-            )
+            transform_list.append(transforms.Normalize(mean=self.mean, std=self.std))
         return transforms.Compose(transform_list)
 
-    def _load_single_image(
-        self, path: Union[str, Path]
-    ) -> torch.Tensor:
+    def _zero_image_tensor(self) -> torch.Tensor:
+        """Return a zero tensor matching the expected image shape (C, H, W).
+
+        Used as a placeholder when an image path is an empty string.
+        Channel count is inferred from self.n_channels if available,
+        otherwise derived from self.mode ("L"→1, "RGBA"→4, else 3).
+
+        Returns:
+            Zero tensor of shape (C, image_size, image_size).
+        """
+        if self.n_channels is not None:
+            c = self.n_channels
+        elif self.mode == "L":
+            c = 1
+        elif self.mode == "RGBA":
+            c = 4
+        else:
+            c = 3
+        return torch.zeros(c, self.image_size, self.image_size)
+
+    def _load_single_image(self, path: Union[str, Path]) -> torch.Tensor:
         """Load and transform a single image from disk.
 
+        If path equals missing_path_token, returns a zero tensor of
+        the same shape as a normal image (C, H, W) via _zero_image_tensor.
+
         Called internally by process() for each image path in
         the input list.
 
         Args:
-            path: Path to the image file.
+            path: Path to the image file. If this equals
+                missing_path_token, a zero-filled placeholder tensor
+                is returned instead.
 
         Returns:
             Transformed image tensor of shape (C, H, W).
@@ -183,18 +198,16 @@ def _load_single_image(
         Raises:
             FileNotFoundError: If the image file does not exist.
         """
+        if self.padding is not None and str(path) == self.padding:
+            return self._zero_image_tensor()
         image_path = Path(path)
         if not image_path.exists():
-            raise FileNotFoundError(
-                f"Image file not found: {image_path}"
-            )
+            raise FileNotFoundError(f"Image file not found: {image_path}")
         with Image.open(image_path) as img:
             img.load()
             return self.transform(img)
 
-    def fit(
-        self, samples: Iterable[Dict[str, Any]], field: str
-    ) -> None:
+    def fit(self, samples: Iterable[Dict[str, Any]], field: str) -> None:
         """Fit the processor by inferring n_channels from data.
 
         Scans samples to find the first valid entry for the given
@@ -214,8 +227,10 @@ def fit(
             for sample in samples:
                 if field in sample and sample[field] is not None:
                     image_paths, _ = sample[field]
-                    if len(image_paths) > 0:
-                        path = Path(image_paths[0])
+                    for raw_path in image_paths:
+                        if self.padding is not None and str(raw_path) == self.padding:
+                            continue
+                        path = Path(raw_path)
                         if path.exists():
                             with Image.open(path) as img:
                                 if img.mode == "L":
@@ -225,14 +240,14 @@ def fit(
                                 else:
                                     self.n_channels = 3
                             break
+                if self.n_channels is not None:
+                    break
             if self.n_channels is None:
                 self.n_channels = 3
 
     def process(
         self,
-        value: Tuple[
-            List[Union[str, Path]], List[float]
-        ],
+        value: Tuple[List[Union[str, Path]], List[float]],
     ) -> Tuple[torch.Tensor, torch.Tensor, str]:
         """Process paired image paths and timestamps.
 
@@ -278,15 +293,10 @@ def process(
         if len(image_paths) == 0:
             raise ValueError("image_paths must be non-empty.")
 
-        paired = sorted(
-            zip(time_diffs, image_paths), key=lambda x: x[0]
-        )
+        paired = sorted(zip(time_diffs, image_paths), key=lambda x: x[0])
 
-        if (
-            self.max_images is not None
-            and len(paired) > self.max_images
-        ):
-            paired = paired[-self.max_images:]
+        if self.max_images is not None and len(paired) > self.max_images:
+            paired = paired[-self.max_images :]
 
         timestamps = []
         image_tensors = []
@@ -295,9 +305,7 @@ def process(
             timestamps.append(t)
 
         images = torch.stack(image_tensors, dim=0)
-        timestamps = torch.tensor(
-            timestamps, dtype=torch.float32
-        )
+        timestamps = torch.tensor(timestamps, dtype=torch.float32)
 
         if self.n_channels is None:
             self.n_channels = images.shape[1]
@@ -344,5 +352,6 @@ def __repr__(self) -> str:
             f"mean={self.mean}, "
             f"std={self.std}, "
             f"mode={self.mode}, "
-            f"max_images={self.max_images})"
-        )
\ No newline at end of file
+            f"max_images={self.max_images}, "
+            f"padding={self.padding!r})"
+        )
diff --git a/pyhealth/processors/tuple_time_text_processor.py b/pyhealth/processors/tuple_time_text_processor.py
index bbe74c4e6..28b21a1c5 100644
--- a/pyhealth/processors/tuple_time_text_processor.py
+++ b/pyhealth/processors/tuple_time_text_processor.py
@@ -5,6 +5,7 @@
 from . import register_processor
 
 logger = logging.getLogger(__name__)
+_MISSING_TEXT_TOKEN = "[MISSING_TEXT]"
 
 @register_processor("tuple_time_text")
 class TupleTimeTextProcessor(TemporalFeatureProcessor):
@@ -81,6 +82,41 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str]
                     - str: Type tag
         """
         texts, time_diffs = value
+        texts = list(texts or [])
+        time_diffs = list(time_diffs or [])
+
+        # Keep text/time aligned and filter malformed text entries.
+        pair_count = min(len(texts), len(time_diffs))
+        cleaned_texts: List[str] = []
+        cleaned_times: List[float] = []
+        for i in range(pair_count):
+            raw_text = texts[i]
+            raw_time = time_diffs[i]
+
+            # Normalize text; skip null/whitespace-only entries.
+            if raw_text is None:
+                continue
+            text = str(raw_text).strip()
+            if text == "":
+                continue
+
+            # Best-effort float normalization; skip unparseable timestamps.
+            try:
+                t = float(raw_time)
+            except (TypeError, ValueError):
+                continue
+
+            cleaned_texts.append(text)
+            cleaned_times.append(t)
+
+        # Fast tokenizer path crashes on empty batches; force a single
+        # missingness token when all notes are empty/malformed.
+        if len(cleaned_texts) == 0:
+            cleaned_texts = [_MISSING_TEXT_TOKEN]
+            cleaned_times = [0.0]
+
+        texts = cleaned_texts
+        time_diffs = cleaned_times
         time_tensor = torch.tensor(time_diffs, dtype=torch.float32)
 
         if self.tokenizer is not None:
@@ -93,16 +129,18 @@ def process(self, value: Tuple[List[str], List[float]]) -> Union[Tuple[List[str]
                 return_tensors="pt"
             )
             
-            input_ids = encoded["input_ids"]
-            attention_mask = encoded["attention_mask"]
+            input_ids = encoded.get("input_ids")
+            if input_ids is None:
+                raise ValueError("Tokenizer output is missing required `input_ids`.")
+
+            attention_mask = encoded.get("attention_mask")
+            if attention_mask is None:
+                attention_mask = torch.ones_like(input_ids)
             
             # Not all tokenizers return token_type_ids (e.g. RoBERTa might not, BERT does)
-            if "token_type_ids" in encoded:
-                token_type_ids = encoded["token_type_ids"]
-            else:
-                # meaningful text usually 0, padding 0? BERT uses 0 for sent A. 
-                # If not provided, we can just use zeros or omit. 
-                # For consistency with schema, let's provide zeros if expected.
+            token_type_ids = encoded.get("token_type_ids")
+            if token_type_ids is None:
+                # Some tokenizers do not return token_type_ids.
                 token_type_ids = torch.zeros_like(input_ids)
 
             return input_ids, attention_mask, token_type_ids, time_tensor, self.type_tag
diff --git a/pyhealth/tasks/__init__.py b/pyhealth/tasks/__init__.py
index 797988377..4b265eef9 100644
--- a/pyhealth/tasks/__init__.py
+++ b/pyhealth/tasks/__init__.py
@@ -44,6 +44,13 @@
 from .mortality_prediction_stagenet_mimic4 import (
     MortalityPredictionStageNetMIMIC4,
 )
+from .multimodal_mimic4 import (
+    ICDLabsMIMIC4,
+    LabsMIMIC4,
+    NotesLabsMIMIC4,
+    NotesLabsCXRMIMIC4,
+    CXRMIMIC4,
+)
 from .patient_linkage import patient_linkage_mimic3_fn
 from .readmission_prediction import (
     ReadmissionPredictionEICU,
diff --git a/pyhealth/tasks/multimodal_mimic4.py b/pyhealth/tasks/multimodal_mimic4.py
new file mode 100644
index 000000000..033c22d6c
--- /dev/null
+++ b/pyhealth/tasks/multimodal_mimic4.py
@@ -0,0 +1,1241 @@
+import logging
+import re
+from datetime import datetime, timedelta
+from typing import Any, Dict, List, Optional, Union, Tuple, ClassVar
+
+from pyhealth.tasks.base_task import BaseTask
+
+logger = logging.getLogger(__name__)
+
+
+class BaseMultimodalMIMIC4Task(BaseTask):
+    """Base class for multimodal MIMIC-IV tasks.
+
+    Provides shared constants and utility methods used across all multimodal
+    task variants (notes, ICD codes, lab values).
+    """
+
+    MISSING_TEXT_TOKEN: ClassVar[str] = ""
+    MISSING_CODE_TOKEN: ClassVar[str] = "<missing>"
+    MISSING_FLOAT_TOKEN: ClassVar[float] = 0.0
+
+    LAB_CATEGORIES: ClassVar[Dict[str, List[str]]] = {
+        "Sodium": ["50824", "52455", "50983", "52623"],
+        "Potassium": ["50822", "52452", "50971", "52610"],
+        "Chloride": ["50806", "52434", "50902", "52535"],
+        "Bicarbonate": ["50803", "50804"],
+        "Glucose": ["50809", "52027", "50931", "52569"],
+        "Calcium": ["50808", "51624"],
+        "Magnesium": ["50960"],
+        "Anion Gap": ["50868", "52500"],
+        "Osmolality": ["52031", "50964", "51701"],
+        "Phosphate": ["50970"],
+    }
+
+    LAB_CATEGORY_NAMES: ClassVar[List[str]] = [
+        "Sodium",
+        "Potassium",
+        "Chloride",
+        "Bicarbonate",
+        "Glucose",
+        "Calcium",
+        "Magnesium",
+        "Anion Gap",
+        "Osmolality",
+        "Phosphate",
+    ]
+
+    LABITEMS: ClassVar[List[str]] = [
+        item for itemids in LAB_CATEGORIES.values() for item in itemids
+    ]
+
+    VITAL_CATEGORIES: ClassVar[Dict[str, List[str]]] = {
+        "HeartRate": ["220045", "220046", "220047"],
+        "SysBP": ["220050", "220179"],
+        "DiasBP": ["220051", "220180"],
+        "MeanBP": ["220052", "220181"],
+        "RespRate": ["220210", "220227"],
+        "SpO2": ["220277"],
+        "Temperature": ["223761", "223762"],
+    }
+
+    VITAL_CATEGORY_NAMES: ClassVar[List[str]] = [
+        "HeartRate",
+        "SysBP",
+        "DiasBP",
+        "MeanBP",
+        "RespRate",
+        "SpO2",
+        "Temperature",
+    ]
+
+    VITALITEMS: ClassVar[List[str]] = [
+        item for itemids in VITAL_CATEGORIES.values() for item in itemids
+    ]
+
+    RADIOLOGY_CLINICAL_HEADERS: ClassVar[List[str]] = [
+        "indication",
+        "impression",
+        # "findings",
+        # "clinical history",
+        # "history",
+        # "comparison",
+        # "technique",
+        # "conclusion",
+        # "summary"
+    ]
+
+    DISCHARGE_CLINICAL_HEADERS: ClassVar[List[str]] = [
+        "chief complaint",
+        # "history of present illness",
+        # "hpi",
+        # "past medical history",
+        # "past medical and surgical history",
+        # "past medical/surgical history",
+        # "past surgical history",
+        # "medications on admission",
+        # "admission medications",
+        # "home medications",
+        # "social history",
+        # "family history",
+        # "allergies",
+        # "review of systems",
+    ]
+
+    def __init__(
+        self,
+        window_hours: Optional[float] = None,
+    ):
+        self.window_hours = window_hours
+
+    @staticmethod
+    def _clean_text(text: Optional[str]) -> Optional[str]:
+        """Return text if non-empty, otherwise None."""
+        return text if text else None
+
+    @staticmethod
+    def _parse_note_sections(text: str, note_type: str) -> Dict[str, str]:
+        """Split a note into {lowercased_header: content_text} pairs."""
+        ext_text = text + '\n\n'
+        if note_type == "radiology":
+            section_re = re.compile(r'([a-zA-Z ]+):[ \t\n]+(.+?)\n{2,}', re.DOTALL)
+        elif note_type == "discharge":
+            section_re = re.compile(r'([a-zA-Z ]+):\n+(.+?)\n{2,}', re.DOTALL)
+        else:
+            raise ValueError(f"Note Type '{note_type}' not supported.")
+        return {
+            m.group(1).strip().lower(): m.group(2).strip()
+            for m in section_re.finditer(ext_text)
+            if m.end() - m.start() > 0
+        }
+
+    @staticmethod
+    def _parse_datetime(value: Any) -> Optional[datetime]:
+        if isinstance(value, datetime):
+            return value
+        if isinstance(value, str):
+            for fmt in ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d"):
+                try:
+                    return datetime.strptime(value, fmt)
+                except ValueError:
+                    continue
+        return None
+
+    @staticmethod
+    def _to_hours(delta_seconds: float) -> float:
+        return delta_seconds / 3600.0
+
+    def _compute_effective_window(
+        self,
+        admissions_to_process: List[Any],
+    ) -> Tuple[datetime, Optional[datetime]]:
+        """Compute effective start/end from the global span of processed admissions.
+
+        Returns:
+            Tuple of (effective_start, effective_end).
+        """
+        global_start = admissions_to_process[0].timestamp
+        global_end: Optional[datetime] = None
+
+        for a in admissions_to_process:
+            dt = self._parse_datetime(getattr(a, "dischtime", None))
+            if dt is not None and (global_end is None or dt > global_end):
+                global_end = dt
+
+        if self.window_hours is not None:
+            effective_start = global_start
+            effective_end = effective_start + timedelta(hours=self.window_hours)
+            return effective_start, effective_end
+
+        effective_start = global_start
+        effective_end = global_end
+
+        return effective_start, effective_end
+
+    def _build_admissions_to_process(self, patient: Any) -> Tuple[List[Any], int]:
+        """Build admissions to process and derive mortality label.
+
+        Includes all admissions up to and including the first death admission.
+        Patients who die in their first (and only) admission are included as
+        positives — previously they were dropped, which silently discarded most
+        ICU mortality positives and collapsed positive rate from ~10% to ~2.7%.
+        This now matches stagenet's semantics: use all available admission data
+        and label as positive if any admission has hospital_expire_flag=1.
+        """
+        admissions = patient.get_events(event_type="admissions")
+        if len(admissions) == 0:
+            return [], 0
+
+        admissions_to_process: List[Any] = []
+        mortality_label = 0
+
+        for admission in admissions:
+            admissions_to_process.append(admission)
+            if admission.hospital_expire_flag in [1, "1"]:
+                mortality_label = 1
+                break
+
+        return admissions_to_process, mortality_label
+
+    def _collect_icd_codes(self, patient: Any, hadm_id: Any) -> List[str]:
+        """Collect ICD diagnosis and procedure codes for one admission.
+
+        Returns:
+            List of ICD code strings, or an empty list if none found.
+        """
+        diagnoses_icd = patient.get_events(
+            event_type="diagnoses_icd", filters=[("hadm_id", "==", hadm_id)]
+        )
+        procedures_icd = patient.get_events(
+            event_type="procedures_icd", filters=[("hadm_id", "==", hadm_id)]
+        )
+        return [
+            e.icd_code for e in diagnoses_icd if hasattr(e, "icd_code") and e.icd_code
+        ] + [
+            e.icd_code for e in procedures_icd if hasattr(e, "icd_code") and e.icd_code
+        ]
+
+    def _collect_labs(
+        self,
+        patient: Any,
+        admission_time: datetime,
+        end_time: datetime,
+    ) -> Tuple[List[float], List[List[float]], List[List[bool]]]:
+        """Collect lab values and observation masks for one admission.
+
+        Args:
+            patient: Patient object.
+            admission_time: Start of the window; times are relative to this.
+            end_time: End of the window (inclusive).
+
+        Returns:
+            Tuple of (lab_times, lab_values, lab_masks). ``lab_masks`` is a
+            parallel boolean tensor where ``True`` means observed and ``False``
+            means imputed with 0.0. Falls back to a single missing placeholder
+            row when no valid lab events are found.
+        """
+        try:
+            import polars as pl
+        except ImportError as exc:
+            raise ImportError("Polars is required for lab collection.") from exc
+
+        labevents_df = patient.get_events(
+            event_type="labevents",
+            start=admission_time,
+            end=end_time,
+            return_df=True,
+        )
+
+        lab_times: List[float] = []
+        lab_values: List[List[float]] = []
+        lab_masks: List[List[bool]] = []
+
+        labevents_df = labevents_df.filter(
+            pl.col("labevents/itemid").is_in(self.LABITEMS)
+        )
+        if labevents_df.height > 0:
+            labevents_df = labevents_df.with_columns(
+                pl.col("labevents/storetime").str.strptime(
+                    pl.Datetime, "%Y-%m-%d %H:%M:%S"
+                )
+            )
+            labevents_df = labevents_df.filter(
+                pl.col("labevents/storetime") <= end_time
+            )
+            if labevents_df.height > 0:
+                labevents_df = labevents_df.select(
+                    pl.col("timestamp"),
+                    pl.col("labevents/itemid"),
+                    pl.col("labevents/valuenum").cast(pl.Float64),
+                )
+                for lab_ts in sorted(labevents_df["timestamp"].unique().to_list()):
+                    ts_labs = labevents_df.filter(pl.col("timestamp") == lab_ts)
+                    lab_vector: List[float] = []
+                    lab_mask: List[bool] = []
+                    for category_name in self.LAB_CATEGORY_NAMES:
+                        category_value = self.MISSING_FLOAT_TOKEN
+                        observed = False
+                        for itemid in self.LAB_CATEGORIES[category_name]:
+                            matching = ts_labs.filter(
+                                pl.col("labevents/itemid") == itemid
+                            )
+                            if matching.height > 0:
+                                category_value = matching["labevents/valuenum"][0]
+                                observed = True
+                                break
+                        lab_vector.append(category_value)
+                        lab_mask.append(observed)
+                    lab_times.append(
+                        self._to_hours((lab_ts - admission_time).total_seconds())
+                    )
+                    lab_values.append(lab_vector)
+                    lab_masks.append(lab_mask)
+            else:  # If missing lab for a given admission
+                lab_values.append(
+                    [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)
+                )
+                lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+                lab_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if len(lab_values) == 0:  # If missing lab for ALL admissions
+            lab_values.append([self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES))
+            lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+            lab_times.append(self.MISSING_FLOAT_TOKEN)
+        return lab_times, lab_values, lab_masks
+
+    def _collect_vitals(
+        self,
+        patient: Any,
+        admission_time: datetime,
+        end_time: datetime,
+    ) -> Tuple[List[float], List[List[float]], List[List[bool]]]:
+        """Collect vital sign values and observation masks for one admission.
+
+        Returns:
+            Tuple of (vital_times, vital_values, vital_masks).
+            vital_masks is parallel boolean: True = observed, False = imputed 0.0.
+            Falls back to a single missing-placeholder row when no valid vitals are found.
+        """
+        try:
+            import polars as pl
+        except ImportError as exc:
+            raise ImportError("Polars is required for vitals collection.") from exc
+
+        chartevents_df = patient.get_events(
+            event_type="chartevents",
+            start=admission_time,
+            end=end_time,
+            return_df=True,
+        )
+
+        vital_times: List[float] = []
+        vital_values: List[List[float]] = []
+        vital_masks: List[List[bool]] = []
+
+        chartevents_df = chartevents_df.filter(
+            pl.col("chartevents/itemid").is_in(self.VITALITEMS)
+        )
+        if chartevents_df.height > 0:
+            chartevents_df = chartevents_df.with_columns(
+                pl.col("chartevents/storetime").str.strptime(
+                    pl.Datetime, "%Y-%m-%d %H:%M:%S"
+                )
+            )
+            chartevents_df = chartevents_df.filter(
+                pl.col("chartevents/storetime") <= end_time
+            )
+            if chartevents_df.height > 0:
+                chartevents_df = chartevents_df.select(
+                    pl.col("timestamp"),
+                    pl.col("chartevents/itemid"),
+                    pl.col("chartevents/valuenum").cast(pl.Float64),
+                )
+                for vital_ts in sorted(chartevents_df["timestamp"].unique().to_list()):
+                    ts_vitals = chartevents_df.filter(pl.col("timestamp") == vital_ts)
+                    vital_vector: List[float] = []
+                    vital_mask: List[bool] = []
+                    for category_name in self.VITAL_CATEGORY_NAMES:
+                        category_value = self.MISSING_FLOAT_TOKEN
+                        observed = False
+                        for itemid in self.VITAL_CATEGORIES[category_name]:
+                            matching = ts_vitals.filter(
+                                pl.col("chartevents/itemid") == itemid
+                            )
+                            if matching.height > 0:
+                                category_value = matching["chartevents/valuenum"][0]
+                                observed = True
+                                break
+                        vital_vector.append(category_value)
+                        vital_mask.append(observed)
+                    vital_times.append(
+                        self._to_hours((vital_ts - admission_time).total_seconds())
+                    )
+                    vital_values.append(vital_vector)
+                    vital_masks.append(vital_mask)
+
+        if len(vital_values) == 0:
+            vital_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.VITAL_CATEGORY_NAMES)
+            )
+            vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES))
+            vital_times.append(self.MISSING_FLOAT_TOKEN)
+
+        return vital_times, vital_values, vital_masks
+
+    def _collect_notes(
+        self,
+        patient: Any,
+        note_event_type: str,
+        hadm_id: Any,
+        admission_time: datetime,
+        start_time: Optional[datetime] = None,
+        end_time: Optional[datetime] = None,
+        section_headers: Optional[List[str]] = None,
+        fallback_to_full_note: bool = False,
+    ) -> Tuple[List[str], List[float]]:
+        """Collect notes of a given type for one admission.
+
+        Args:
+            patient: Patient object.
+            note_event_type: Event type string (e.g. "discharge", "radiology").
+            hadm_id: Admission ID to filter by.
+            admission_time: Admission start time; used to compute time offsets.
+            start_time: Optional start of the time window.
+            end_time: Optional end of the time window.
+            section_headers: When provided, extract only these named sections
+                from each note (lowercased match against parsed headers).
+            fallback_to_full_note: When True (default), falls back to the full
+                note text if no matching sections are found. When False, notes
+                with no matching sections are dropped entirely.
+
+        Returns:
+            Tuple of (texts, hours_from_admission). Falls back to
+            ``([MISSING_TEXT_TOKEN], [MISSING_FLOAT_TOKEN])`` when the events
+            list is empty.
+        """
+        notes = patient.get_events(
+            event_type=note_event_type,
+            start=start_time,
+            end=end_time,
+            filters=[("hadm_id", "==", hadm_id)],
+        )
+
+        texts: List[str] = []
+        note_times: List[float] = []
+        for note in notes:
+            try:
+                note_text = self._clean_text(note.text)
+                if note_text:
+                    if section_headers is not None:
+                        parsed = self._parse_note_sections(note_text, note_type=note_event_type)
+                        extracted = [f"{k}: {v}" for k, v in parsed.items() if k in section_headers and v]
+                        if extracted:
+                            note_text = " [SEP] ".join(extracted)
+                        elif not fallback_to_full_note:
+                            continue
+
+                    time_from_admission = self._to_hours(
+                        (note.timestamp - admission_time).total_seconds()
+                    )
+                    texts.append(note_text)
+                    note_times.append(time_from_admission)
+            except (
+                AttributeError
+            ):  # note object is missing .text or .timestamp attribute (e.g. malformed note)
+                pass
+
+        return texts, note_times
+
+
+class ICDLabsMIMIC4(BaseMultimodalMIMIC4Task):
+    """Task for ICD codes + lab values mortality prediction using MIMIC-IV.
+
+    A notes-free structured-EHR task that uses only:
+
+    - **ICD codes**: diagnosis and procedure codes per admission, processed by
+      ``StageNetProcessor`` with inter-admission time offsets.
+    - **Lab values**: 10-dimensional lab vectors (one per lab category) at each
+      measurement timestamp, processed by ``StageNetTensorProcessor``.
+
+    Examples:
+        >>> from pyhealth.datasets import MIMIC4Dataset
+        >>> from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4
+        >>> dataset = MIMIC4Dataset(
+        ...     ehr_root="/path/to/mimic-iv/2.2",
+        ...     ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"],
+        ... )
+        >>> task = ICDLabsMIMIC4()
+        >>> samples = dataset.set_task(task)
+    """
+
+    PADDING: int = 0
+
+    task_name: str = "ICDLabsMIMIC4"
+    input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = {
+        "icd_codes": ("stagenet", {"padding": PADDING}),
+        "labs": ("stagenet_tensor", {}),
+        "labs_mask": ("stagenet_tensor", {}),
+    }
+    output_schema: Dict[str, str] = {"mortality": "binary"}
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:
+        demographics = patient.get_events(event_type="patients")
+        if not demographics:
+            return []
+
+        admissions_to_process, mortality_label = self._build_admissions_to_process(
+            patient
+        )
+
+        if len(admissions_to_process) == 0:
+            return []
+
+        effective_start, effective_end = self._compute_effective_window(
+            admissions_to_process
+        )
+
+        all_icd_codes: List[List[str]] = []
+        all_icd_times: List[float] = []
+        all_lab_values: List[List[float]] = []
+        all_lab_masks: List[List[bool]] = []
+        all_lab_times: List[float] = []
+        previous_admission_time = None
+
+        for admission in admissions_to_process:
+            admission_time = admission.timestamp
+
+            try:
+                admission_dischtime = datetime.strptime(
+                    admission.dischtime, "%Y-%m-%d %H:%M:%S"
+                )
+            except (ValueError, AttributeError):
+                admission_dischtime = admission_time
+            if admission_dischtime < admission_time:
+                admission_dischtime = admission_time
+
+            visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id)
+            if visit_icd_codes:
+                if previous_admission_time is None:
+                    time_from_previous = 0.0
+                else:
+                    time_from_previous = self._to_hours(
+                        (admission_time - previous_admission_time).total_seconds()
+                    )
+                all_icd_codes.append(visit_icd_codes)
+                all_icd_times.append(time_from_previous)
+            else:
+                all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+                all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+
+            previous_admission_time = admission_time
+
+            lab_times, lab_values, lab_masks = self._collect_labs(
+                patient=patient,
+                admission_time=admission_time,
+                end_time=admission_dischtime,
+            )
+            all_lab_times.extend(lab_times)
+            all_lab_values.extend(lab_values)
+            all_lab_masks.extend(lab_masks)
+
+        if len(all_lab_values) == 0:
+            all_lab_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)
+            )
+            all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+            all_lab_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if len(all_icd_codes) == 0:
+            all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+            all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+
+        single_patient_longitudinal_record = {
+            "patient_id": patient.patient_id,
+            "icd_codes": (all_icd_times, all_icd_codes),
+            "labs": (all_lab_times, all_lab_values),
+            "labs_mask": (all_lab_times, all_lab_masks),
+            "mortality": mortality_label,
+            "window_start": effective_start,
+            "window_end": effective_end,
+        }
+
+        return [single_patient_longitudinal_record]
+
+
+class NotesLabsMIMIC4(BaseMultimodalMIMIC4Task):
+    """Mortality prediction from admission-context notes and lab values.
+
+    Follows the approach of Lee et al. (2023): use text that is clinically
+    available *at admission* rather than discharge notes. ICD codes and vitals
+    are excluded by default but can be re-enabled for ablation experiments.
+
+    Text is extracted from the MIMIC-IV discharge note by parsing the Chief
+    Complaint, History of Present Illness, Past Medical History, and Medications
+    on Admission sections — all of which describe the patient's state at the
+    start of the stay. The extracted text is assigned timestamp 0.0.
+
+    Radiology reports are also included, parsed for their Indication and
+    Impression sections and bounded to the same observation window as labs
+    (rather than timestamp 0.0), since — unlike the discharge summary — they
+    are written at exam time and describe findings from later in the stay.
+
+    Fields:
+        admission_note_times: Admission-context discharge-note text at time
+            0.0, plus in-window radiology note text at its exam-relative
+            timestamp.
+        labs: 10-dim lab vectors at each measurement timestamp.
+        labs_mask: Boolean observation mask parallel to ``labs``.
+        vitals: (only when ``include_vitals=True``) 7-dim vital-sign vectors at
+            each measurement timestamp.
+        vitals_mask: (only when ``include_vitals=True``) Boolean observation mask
+            parallel to ``vitals``.
+        icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes
+            per admission with inter-admission time offsets.
+
+    Args:
+        window_hours: Hours from admission for lab/vital collection. ``None``
+            collects for the full admission span. Default: 24.
+        include_icd: When ``True``, collect discharge-coded ICD codes and add
+            ``icd_codes`` to the sample dict / input schema. Default: ``False``.
+        include_vitals: When ``True``, collect ICU vital signs from chartevents
+            and add ``vitals`` and ``vitals_mask`` to the sample dict / input
+            schema. Default: ``False``.
+    """
+
+    PADDING: int = 0
+
+    task_name: str = "NotesLabsMIMIC4"
+
+    _BASE_INPUT_SCHEMA: ClassVar[Dict] = {
+        "admission_note_times": (
+            "tuple_time_text",
+            {
+                "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT",
+                "type_tag": "note",
+            },
+        ),
+        "labs": ("stagenet_tensor", {}),
+        "labs_mask": ("stagenet_tensor", {}),
+    }
+
+    input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA
+    output_schema: Dict[str, str] = {"mortality": "binary"}
+
+    def __init__(
+        self,
+        window_hours: Optional[float] = None,
+        include_icd: bool = False,
+        include_vitals: bool = False,
+    ) -> None:
+        super().__init__(window_hours=window_hours)
+        self.include_icd = include_icd
+        self.include_vitals = include_vitals
+        schema = dict(self._BASE_INPUT_SCHEMA)
+        if include_vitals:
+            schema["vitals"] = ("stagenet_tensor", {})
+            schema["vitals_mask"] = ("stagenet_tensor", {})
+        if include_icd:
+            schema["icd_codes"] = ("stagenet", {"padding": self.PADDING})
+        self.input_schema = schema
+        logger.info(
+            "NotesLabsMIMIC4: filtering discharge notes to sections: %s; "
+            "radiology notes to sections: %s",
+            self.DISCHARGE_CLINICAL_HEADERS,
+            self.RADIOLOGY_CLINICAL_HEADERS,
+        )
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:
+        if not patient.get_events(event_type="patients"):
+            return []
+
+        admissions_to_process, mortality_label = self._build_admissions_to_process(
+            patient
+        )
+        if not admissions_to_process:
+            return []
+
+        effective_start, effective_end = self._compute_effective_window(
+            admissions_to_process
+        )
+
+        all_note_texts: List[str] = []
+        all_note_times: List[float] = []
+        all_lab_values: List[List[float]] = []
+        all_lab_masks: List[List[bool]] = []
+        all_lab_times: List[float] = []
+        all_vital_values: List[List[float]] = []
+        all_vital_masks: List[List[bool]] = []
+        all_vital_times: List[float] = []
+        all_icd_codes: List[List[str]] = []
+        all_icd_times: List[float] = []
+        previous_admission_time = None
+
+        for admission in admissions_to_process:
+            admission_time = admission.timestamp
+
+            try:
+                admission_dischtime = datetime.strptime(
+                    admission.dischtime, "%Y-%m-%d %H:%M:%S"
+                )
+            except (ValueError, AttributeError):
+                admission_dischtime = admission_time
+            if admission_dischtime < admission_time:
+                admission_dischtime = admission_time
+
+            note_texts, note_times = self._collect_notes(
+                patient,
+                "discharge",
+                admission.hadm_id,
+                admission_time,
+                section_headers=self.DISCHARGE_CLINICAL_HEADERS,
+            )
+            all_note_texts.extend(note_texts)
+            all_note_times.extend(note_times)
+
+            # Labs within the observation window
+            lab_end = (
+                effective_end
+                if self.window_hours is not None
+                else admission_dischtime
+            )
+            lab_times, lab_values, lab_masks = self._collect_labs(
+                patient=patient,
+                admission_time=admission_time,
+                end_time=lab_end,
+            )
+            all_lab_times.extend(lab_times)
+            all_lab_values.extend(lab_values)
+            all_lab_masks.extend(lab_masks)
+
+            # Radiology notes within the observation window. Unlike the
+            # discharge note (a retrospective summary parsed for its
+            # admission-context sections), radiology reports are written at
+            # exam time, so they're bounded to the same window as labs/vitals
+            # to avoid pulling in findings from later in the stay.
+            radiology_texts, radiology_times = self._collect_notes(
+                patient,
+                "radiology",
+                admission.hadm_id,
+                admission_time,
+                start_time=admission_time,
+                end_time=lab_end,
+                section_headers=self.RADIOLOGY_CLINICAL_HEADERS,
+            )
+            all_note_texts.extend(radiology_texts)
+            all_note_times.extend(radiology_times)
+
+            # Vitals within the observation window
+            if self.include_vitals:
+                vital_times, vital_values, vital_masks = self._collect_vitals(
+                    patient=patient,
+                    admission_time=admission_time,
+                    end_time=lab_end,
+                )
+                all_vital_times.extend(vital_times)
+                all_vital_values.extend(vital_values)
+                all_vital_masks.extend(vital_masks)
+
+            if self.include_icd:
+                visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id)
+                time_from_previous = (
+                    0.0
+                    if previous_admission_time is None
+                    else self._to_hours(
+                        (admission_time - previous_admission_time).total_seconds()
+                    )
+                )
+                if visit_icd_codes:
+                    all_icd_codes.append(visit_icd_codes)
+                    all_icd_times.append(time_from_previous)
+                else:
+                    all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+                    all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+                previous_admission_time = admission_time
+
+        if not all_lab_values:
+            all_lab_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)
+            )
+            all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+            all_lab_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if self.include_vitals and not all_vital_values:
+            all_vital_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.VITAL_CATEGORY_NAMES)
+            )
+            all_vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES))
+            all_vital_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if not all_note_texts:
+            all_note_texts = [self.MISSING_TEXT_TOKEN]
+            all_note_times = [self.MISSING_FLOAT_TOKEN]
+
+        record: Dict[str, Any] = {
+            "patient_id": patient.patient_id,
+            "admission_note_times": (all_note_texts, all_note_times),
+            "labs": (all_lab_times, all_lab_values),
+            "labs_mask": (all_lab_times, all_lab_masks),
+            "mortality": mortality_label,
+            "window_start": effective_start,
+            "window_end": effective_end,
+        }
+
+        if self.include_vitals:
+            record["vitals"] = (all_vital_times, all_vital_values)
+            record["vitals_mask"] = (all_vital_times, all_vital_masks)
+
+        if self.include_icd:
+            if not all_icd_codes:
+                all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+                all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+            record["icd_codes"] = (all_icd_times, all_icd_codes)
+
+        return [record]
+
+
+class NotesLabsCXRMIMIC4(BaseMultimodalMIMIC4Task):
+    """Mortality prediction from admission-context notes, labs, and CXR.
+
+    Extends ``NotesLabsMIMIC4`` with chest X-ray images: the same
+    admission-context discharge-note sections, in-window radiology reports,
+    and labs, plus CXR studies (StudyDate+StudyTime from the ``metadata``
+    event table) bounded to the same observation window as labs/radiology.
+    ICD codes and vitals are excluded by default but can be re-enabled for
+    ablation experiments, same as ``NotesLabsMIMIC4``.
+
+    Fields:
+        admission_note_times: Admission-context discharge-note text at time
+            0.0, plus in-window radiology note text at its exam-relative
+            timestamp.
+        labs: 10-dim lab vectors at each measurement timestamp.
+        labs_mask: Boolean observation mask parallel to ``labs``.
+        cxr_image_times: In-window CXR image paths at their exam-relative
+            timestamp.
+        vitals: (only when ``include_vitals=True``) 7-dim vital-sign vectors at
+            each measurement timestamp.
+        vitals_mask: (only when ``include_vitals=True``) Boolean observation mask
+            parallel to ``vitals``.
+        icd_codes: (only when ``include_icd=True``) Diagnosis + procedure codes
+            per admission with inter-admission time offsets.
+
+    Args:
+        window_hours: Hours from admission for lab/vital/CXR collection.
+            ``None`` (default) collects for the full admission span.
+        include_icd: When ``True``, collect discharge-coded ICD codes and add
+            ``icd_codes`` to the sample dict / input schema. Default: ``False``.
+        include_vitals: When ``True``, collect ICU vital signs from chartevents
+            and add ``vitals`` and ``vitals_mask`` to the sample dict / input
+            schema. Default: ``False``.
+    """
+
+    PADDING: int = 0
+
+    task_name: str = "NotesLabsCXRMIMIC4"
+
+    _BASE_INPUT_SCHEMA: ClassVar[Dict] = {
+        "admission_note_times": (
+            "tuple_time_text",
+            {
+                "tokenizer_model": "emilyalsentzer/Bio_ClinicalBERT",
+                "type_tag": "note",
+            },
+        ),
+        "labs": ("stagenet_tensor", {}),
+        "labs_mask": ("stagenet_tensor", {}),
+        "cxr_image_times": (
+            "time_image",
+            {
+                "image_size": 224,
+                "mode": "RGB",
+                "padding": "",
+            },
+        ),
+    }
+
+    input_schema: Dict[str, Union[str, Tuple[str, Dict]]] = _BASE_INPUT_SCHEMA
+    output_schema: Dict[str, str] = {"mortality": "binary"}
+
+    def __init__(
+        self,
+        window_hours: Optional[float] = None,
+        include_icd: bool = False,
+        include_vitals: bool = False,
+    ) -> None:
+        super().__init__(window_hours=window_hours)
+        self.include_icd = include_icd
+        self.include_vitals = include_vitals
+        schema = dict(self._BASE_INPUT_SCHEMA)
+        if include_vitals:
+            schema["vitals"] = ("stagenet_tensor", {})
+            schema["vitals_mask"] = ("stagenet_tensor", {})
+        if include_icd:
+            schema["icd_codes"] = ("stagenet", {"padding": self.PADDING})
+        self.input_schema = schema
+        logger.info(
+            "NotesLabsCXRMIMIC4: filtering discharge notes to sections: %s; "
+            "radiology notes to sections: %s",
+            self.DISCHARGE_CLINICAL_HEADERS,
+            self.RADIOLOGY_CLINICAL_HEADERS,
+        )
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:
+        if not patient.get_events(event_type="patients"):
+            return []
+
+        admissions_to_process, mortality_label = self._build_admissions_to_process(
+            patient
+        )
+        if not admissions_to_process:
+            return []
+
+        effective_start, effective_end = self._compute_effective_window(
+            admissions_to_process
+        )
+
+        all_note_texts: List[str] = []
+        all_note_times: List[float] = []
+        all_lab_values: List[List[float]] = []
+        all_lab_masks: List[List[bool]] = []
+        all_lab_times: List[float] = []
+        all_vital_values: List[List[float]] = []
+        all_vital_masks: List[List[bool]] = []
+        all_vital_times: List[float] = []
+        all_icd_codes: List[List[str]] = []
+        all_icd_times: List[float] = []
+        all_cxr_paths: List[str] = []
+        all_cxr_times: List[float] = []
+        previous_admission_time = None
+
+        for admission in admissions_to_process:
+            admission_time = admission.timestamp
+
+            # Skip admissions that start at or after the observation window
+            # closes, prevents Polars searchsorted OverflowError in CXR lookup.
+            if effective_end is not None and admission_time >= effective_end:
+                continue
+
+            try:
+                admission_dischtime = datetime.strptime(
+                    admission.dischtime, "%Y-%m-%d %H:%M:%S"
+                )
+            except (ValueError, AttributeError):
+                admission_dischtime = admission_time
+            if admission_dischtime < admission_time:
+                admission_dischtime = admission_time
+
+            note_texts, note_times = self._collect_notes(
+                patient,
+                "discharge",
+                admission.hadm_id,
+                admission_time,
+                section_headers=self.DISCHARGE_CLINICAL_HEADERS,
+            )
+            all_note_texts.extend(note_texts)
+            all_note_times.extend(note_times)
+
+            # Labs within the observation window
+            lab_end = (
+                effective_end
+                if self.window_hours is not None
+                else admission_dischtime
+            )
+            lab_times, lab_values, lab_masks = self._collect_labs(
+                patient=patient,
+                admission_time=admission_time,
+                end_time=lab_end,
+            )
+            all_lab_times.extend(lab_times)
+            all_lab_values.extend(lab_values)
+            all_lab_masks.extend(lab_masks)
+
+            # Radiology notes within the observation window. Unlike the
+            # discharge note (a retrospective summary parsed for its
+            # admission-context sections), radiology reports are written at
+            # exam time, so they're bounded to the same window as labs/CXR
+            # to avoid pulling in findings from later in the stay.
+            radiology_texts, radiology_times = self._collect_notes(
+                patient,
+                "radiology",
+                admission.hadm_id,
+                admission_time,
+                start_time=admission_time,
+                end_time=lab_end,
+                section_headers=self.RADIOLOGY_CLINICAL_HEADERS,
+            )
+            all_note_texts.extend(radiology_texts)
+            all_note_times.extend(radiology_times)
+
+            # CXR studies within the same observation window as labs/radiology.
+            # CXR metadata is filtered by timestamp; this includes StudyTime.
+            metadata_events = patient.get_events(
+                event_type="metadata",
+                start=admission_time,
+                end=lab_end,
+            )
+            for event in metadata_events:
+                try:
+                    if event.image_path:
+                        all_cxr_paths.append(event.image_path)
+                        all_cxr_times.append(
+                            self._to_hours(
+                                (event.timestamp - admission_time).total_seconds()
+                            )
+                        )
+                except AttributeError:
+                    continue
+
+            # Vitals within the observation window
+            if self.include_vitals:
+                vital_times, vital_values, vital_masks = self._collect_vitals(
+                    patient=patient,
+                    admission_time=admission_time,
+                    end_time=lab_end,
+                )
+                all_vital_times.extend(vital_times)
+                all_vital_values.extend(vital_values)
+                all_vital_masks.extend(vital_masks)
+
+            if self.include_icd:
+                visit_icd_codes = self._collect_icd_codes(patient, admission.hadm_id)
+                time_from_previous = (
+                    0.0
+                    if previous_admission_time is None
+                    else self._to_hours(
+                        (admission_time - previous_admission_time).total_seconds()
+                    )
+                )
+                if visit_icd_codes:
+                    all_icd_codes.append(visit_icd_codes)
+                    all_icd_times.append(time_from_previous)
+                else:
+                    all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+                    all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+                previous_admission_time = admission_time
+
+        if not all_lab_values:
+            all_lab_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)
+            )
+            all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+            all_lab_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if self.include_vitals and not all_vital_values:
+            all_vital_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.VITAL_CATEGORY_NAMES)
+            )
+            all_vital_masks.append([False] * len(self.VITAL_CATEGORY_NAMES))
+            all_vital_times.append(self.MISSING_FLOAT_TOKEN)
+
+        if not all_note_texts:
+            all_note_texts = [self.MISSING_TEXT_TOKEN]
+            all_note_times = [self.MISSING_FLOAT_TOKEN]
+
+        # time_image processor expects at least one path/time pair.
+        if len(all_cxr_paths) == 0:
+            all_cxr_paths = [self.MISSING_TEXT_TOKEN]
+            all_cxr_times = [self.MISSING_FLOAT_TOKEN]
+
+        record: Dict[str, Any] = {
+            "patient_id": patient.patient_id,
+            "admission_note_times": (all_note_texts, all_note_times),
+            "labs": (all_lab_times, all_lab_values),
+            "labs_mask": (all_lab_times, all_lab_masks),
+            "cxr_image_times": (all_cxr_paths, all_cxr_times),
+            "mortality": mortality_label,
+            "window_start": effective_start,
+            "window_end": effective_end,
+        }
+
+        if self.include_vitals:
+            record["vitals"] = (all_vital_times, all_vital_values)
+            record["vitals_mask"] = (all_vital_times, all_vital_masks)
+
+        if self.include_icd:
+            if not all_icd_codes:
+                all_icd_codes.append([self.MISSING_TEXT_TOKEN])
+                all_icd_times.append(self.MISSING_FLOAT_TOKEN)
+            record["icd_codes"] = (all_icd_times, all_icd_codes)
+
+        return [record]
+
+
+class LabsMIMIC4(BaseMultimodalMIMIC4Task):
+    """EHR-only mortality prediction using lab values — no notes, no ICD codes.
+
+    Serves as the structured-EHR reference baseline for multimodal ablations.
+    Collecting only ``labevents`` keeps the dataset loader fast and avoids any
+    leakage from discharge-coded ICD tables.
+
+    Schema mirrors the ``labs`` / ``labs_mask`` fields from ``NotesLabsMIMIC4``
+    so the same backbone models (MLP, RNN, Transformer, etc.) work unchanged.
+
+    Args:
+        window_hours: Hours from admission to collect lab measurements.
+            ``None`` collects for the full admission span. Default: 24.
+    """
+
+    PADDING: int = 0
+
+    task_name: str = "LabsMIMIC4"
+
+    input_schema: ClassVar[Dict] = {
+        "labs": ("stagenet_tensor", {}),
+        "labs_mask": ("stagenet_tensor", {}),
+    }
+    output_schema: ClassVar[Dict] = {"mortality": "binary"}
+
+    def __init__(self, window_hours: Optional[float] = 24) -> None:
+        super().__init__()
+        self.window_hours = window_hours
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:  # type: ignore[override]
+        admissions_to_process, mortality_label = self._build_admissions_to_process(
+            patient
+        )
+        if not admissions_to_process:
+            return []
+
+        effective_start, effective_end = self._compute_effective_window(
+            admissions_to_process
+        )
+
+        all_lab_times: List[float] = []
+        all_lab_values: List[List[float]] = []
+        all_lab_masks: List[List[bool]] = []
+
+        for admission in admissions_to_process:
+            admission_time = admission.timestamp
+
+            try:
+                admission_dischtime = datetime.strptime(
+                    admission.dischtime, "%Y-%m-%d %H:%M:%S"
+                )
+            except (ValueError, AttributeError):
+                admission_dischtime = admission_time
+            if admission_dischtime < admission_time:
+                admission_dischtime = admission_time
+
+            lab_times, lab_values, lab_masks = self._collect_labs(
+                patient=patient,
+                admission_time=admission_time,
+                end_time=admission_dischtime,
+            )
+            all_lab_times.extend(lab_times)
+            all_lab_values.extend(lab_values)
+            all_lab_masks.extend(lab_masks)
+
+        if len(all_lab_values) == 0:
+            all_lab_values.append(
+                [self.MISSING_FLOAT_TOKEN] * len(self.LAB_CATEGORY_NAMES)
+            )
+            all_lab_masks.append([False] * len(self.LAB_CATEGORY_NAMES))
+            all_lab_times.append(self.MISSING_FLOAT_TOKEN)
+
+        single_patient_longitudinal_record = {
+            "patient_id": patient.patient_id,
+            "labs": (all_lab_times, all_lab_values),
+            "labs_mask": (all_lab_times, all_lab_masks),
+            "mortality": mortality_label,
+            "window_start": effective_start,
+            "window_end": effective_end,
+        }
+
+        return [single_patient_longitudinal_record]
+
+
+class CXRMIMIC4(BaseMultimodalMIMIC4Task):
+    """CXR-only mortality prediction using chest X-ray images.
+
+    Serves as the imaging-only reference baseline for multimodal ablations —
+    no notes, no ICD codes, no labs/vitals — isolating the chest X-ray
+    modality the same way ``LabsMIMIC4`` isolates labs.
+
+    CXR studies are filtered by timestamp (StudyDate+StudyTime, from the
+    ``metadata`` event table) within each admission's observation window.
+
+    Args:
+        window_hours: Hours from admission to collect CXR studies. ``None``
+            collects for the full admission span. Default: None.
+    """
+
+    task_name: str = "CXRMIMIC4"
+
+    input_schema: ClassVar[Dict] = {
+        "cxr_image_times": (
+            "time_image",
+            {
+                "image_size": 224,
+                "mode": "RGB",
+                "padding": "",
+            },
+        ),
+    }
+    output_schema: ClassVar[Dict] = {"mortality": "binary"}
+
+    def __call__(self, patient: Any) -> List[Dict[str, Any]]:  # type: ignore[override]
+        admissions_to_process, mortality_label = self._build_admissions_to_process(
+            patient
+        )
+        if not admissions_to_process:
+            return []
+
+        effective_start, effective_end = self._compute_effective_window(
+            admissions_to_process
+        )
+
+        all_cxr_paths: List[str] = []
+        all_cxr_times: List[float] = []
+
+        for admission in admissions_to_process:
+            admission_time = admission.timestamp
+
+            # Skip admissions that start at or after the observation window
+            # closes, prevents Polars searchsorted OverflowError.
+            if effective_end is not None and admission_time >= effective_end:
+                continue
+
+            try:
+                admission_dischtime = datetime.strptime(
+                    admission.dischtime, "%Y-%m-%d %H:%M:%S"
+                )
+            except (ValueError, AttributeError):
+                admission_dischtime = admission_time
+            if admission_dischtime < admission_time:
+                admission_dischtime = admission_time
+
+            admission_end = admission_dischtime
+            if effective_end is not None and effective_end < admission_end:
+                admission_end = effective_end
+
+            # CXR metadata is filtered by timestamp; this includes StudyTime.
+            metadata_events = patient.get_events(
+                event_type="metadata",
+                start=admission_time,
+                end=admission_end,
+            )
+            for event in metadata_events:
+                try:
+                    if event.image_path:
+                        all_cxr_paths.append(event.image_path)
+                        all_cxr_times.append(
+                            self._to_hours(
+                                (event.timestamp - admission_time).total_seconds()
+                            )
+                        )
+                except AttributeError:
+                    continue
+
+        # time_image processor expects at least one path/time pair.
+        if len(all_cxr_paths) == 0:
+            all_cxr_paths = [self.MISSING_TEXT_TOKEN]
+            all_cxr_times = [self.MISSING_FLOAT_TOKEN]
+
+        single_patient_longitudinal_record = {
+            "patient_id": patient.patient_id,
+            "cxr_image_times": (all_cxr_paths, all_cxr_times),
+            "mortality": mortality_label,
+            "window_start": effective_start,
+            "window_end": effective_end,
+        }
+
+        return [single_patient_longitudinal_record]
diff --git a/pyhealth/trainer.py b/pyhealth/trainer.py
index bc6a28677..ea7f1af70 100644
--- a/pyhealth/trainer.py
+++ b/pyhealth/trainer.py
@@ -1,5 +1,7 @@
+import json
 import logging
 import os
+import time
 from datetime import datetime
 from typing import Callable, Dict, List, Optional, Type
 
@@ -37,6 +39,15 @@ def set_logger(log_path: str) -> None:
     return
 
 
+def _vram_stats(device: str) -> Dict[str, float]:
+    """Returns current and peak VRAM usage in MB for a CUDA device."""
+    if not torch.cuda.is_available() or not str(device).startswith("cuda"):
+        return {}
+    allocated = torch.cuda.memory_allocated(device) / 1024**2
+    peak = torch.cuda.max_memory_allocated(device) / 1024**2
+    return {"vram_allocated_mb": allocated, "vram_peak_mb": peak}
+
+
 def get_metrics_fn(mode: str) -> Callable:
     if mode == "binary":
         return binary_metrics_fn
@@ -126,6 +137,9 @@ def train(
         monitor_criterion: str = "max",
         load_best_model_at_last: bool = True,
         patience=None,
+        accumulation_steps: int = 1,
+        use_amp: bool = False,
+        amp_dtype: str = "bf16",
     ):
         """Trains the model.
 
@@ -145,10 +159,25 @@ def train(
                 Default is True.
             patience: Number of epochs to wait for improvement before early stopping.
                 Default is None, which means no early stopping.
+            accumulation_steps: Gradient accumulation steps to simulate a larger
+                effective batch size. Default is 1 (no accumulation).
+            use_amp: Whether to use automatic mixed precision. Default is False.
+            amp_dtype: AMP dtype — "bf16" (stable, recommended) or "fp16".
+                Default is "bf16".
         """
         if optimizer_params is None:
             optimizer_params = {"lr": 1e-3}
 
+        _amp_dtype = (
+            torch.bfloat16 if amp_dtype == "bf16" else torch.float16
+        )
+        # GradScaler only needed for fp16; bf16 has fp32 dynamic range
+        scaler = (
+            torch.cuda.amp.GradScaler()
+            if (use_amp and _amp_dtype == torch.float16)
+            else None
+        )
+
         # logging
         logger.info("Training:")
         logger.info(f"Batch size: {train_dataloader.batch_size}")
@@ -161,6 +190,8 @@ def train(
         logger.info(f"Monitor criterion: {monitor_criterion}")
         logger.info(f"Epochs: {epochs}")
         logger.info(f"Patience: {patience}")
+        logger.info(f"Accumulation steps: {accumulation_steps}")
+        logger.info(f"AMP: {use_amp} (dtype={amp_dtype})")
 
         # set optimizer
         param = list(self.model.named_parameters())
@@ -184,50 +215,129 @@ def train(
             steps_per_epoch = len(train_dataloader)
         global_step = 0
         patience_counter = 0
+        metrics_history: List[Dict] = []
+        train_start = time.perf_counter()
+        total_skipped_steps = 0
 
         # epoch training loop
-        for epoch in range(epochs):
+        epoch_iterator = tqdm(range(epochs), desc="Epochs", unit="epoch")
+        for epoch in epoch_iterator:
+            epoch_iterator.set_postfix_str(f"{epoch + 1}/{epochs}", refresh=False)
             training_loss = []
+            epoch_skipped_steps = 0
             self.model.zero_grad()
             self.model.train()
+            if torch.cuda.is_available() and str(self.device).startswith("cuda"):
+                torch.cuda.reset_peak_memory_stats(self.device)
+            epoch_start = time.perf_counter()
             # batch training loop
             logger.info("")
-            for _ in trange(
+            for step_idx in trange(
                 steps_per_epoch,
-                desc=f"Epoch {epoch} / {epochs}",
+                desc=f"Epoch {epoch + 1}/{epochs}",
                 smoothing=0.05,
+                leave=False,
             ):
                 try:
                     data = next(data_iterator)
                 except StopIteration:
                     data_iterator = iter(train_dataloader)
                     data = next(data_iterator)
-                # forward
-                output = self.model(**data)
-                loss = output["loss"]
+                # forward (with optional AMP)
+                if use_amp:
+                    with torch.autocast(device_type="cuda", dtype=_amp_dtype):
+                        output = self.model(**data)
+                        loss = output["loss"] / accumulation_steps
+                else:
+                    output = self.model(**data)
+                    loss = output["loss"] / accumulation_steps
                 # backward
-                loss.backward()
-                if max_grad_norm is not None:
-                    torch.nn.utils.clip_grad_norm_(
-                        self.model.parameters(), max_grad_norm
+                if scaler is not None:
+                    scaler.scale(loss).backward()
+                else:
+                    loss.backward()
+                training_loss.append(loss.item() * accumulation_steps)
+                # optimizer step every accumulation_steps batches or epoch end
+                is_update_step = (
+                    (step_idx + 1) % accumulation_steps == 0
+                    or (step_idx + 1) == steps_per_epoch
+                )
+                if is_update_step:
+                    if scaler is not None:
+                        scaler.unscale_(optimizer)
+                    # Always compute the grad norm (even with no clipping
+                    # configured) so non-finite gradients can be detected and
+                    # skipped before they permanently poison the model with
+                    # NaN weights.
+                    grad_norm = torch.nn.utils.clip_grad_norm_(
+                        self.model.parameters(),
+                        max_grad_norm if max_grad_norm is not None else float("inf"),
                     )
-                # update
-                optimizer.step()
-                optimizer.zero_grad()
-                training_loss.append(loss.item())
-                global_step += 1
+                    step_ok = bool(torch.isfinite(grad_norm))
+                    if not step_ok:
+                        epoch_skipped_steps += 1
+                        total_skipped_steps += 1
+                        logger.warning(
+                            f"epoch-{epoch} step-{global_step}: non-finite "
+                            f"gradient norm ({grad_norm}); skipping optimizer "
+                            f"step."
+                        )
+                    if scaler is not None:
+                        if step_ok:
+                            scaler.step(optimizer)
+                        scaler.update()
+                    elif step_ok:
+                        optimizer.step()
+                    optimizer.zero_grad()
+                    global_step += 1
+
+            epoch_time = time.perf_counter() - epoch_start
+            vram = _vram_stats(self.device)
+
+            epochs_done = epoch + 1
+            epochs_left = epochs - epochs_done
+            elapsed_total = time.perf_counter() - train_start
+            avg_epoch_time = elapsed_total / epochs_done
+            eta_s = avg_epoch_time * epochs_left
+            eta_h, eta_rem = divmod(int(eta_s), 3600)
+            eta_m = eta_rem // 60
+            eta_str = f"{eta_h}h{eta_m:02d}m"
+
             # log and save
             logger.info(f"--- Train epoch-{epoch}, step-{global_step} ---")
             logger.info(f"loss: {sum(training_loss) / len(training_loss):.4f}")
+            logger.info(f"epoch_time: {epoch_time:.2f}s  elapsed: {elapsed_total:.0f}s  ETA: {eta_str} ({epochs_done}/{epochs} epochs)")
+            print(f"[ETA] epoch {epochs_done}/{epochs} done in {epoch_time:.0f}s — ETA to finish: {eta_str}", flush=True)
+            if vram:
+                logger.info(
+                    f"vram_peak: {vram['vram_peak_mb']:.1f} MB  "
+                    f"vram_current: {vram['vram_allocated_mb']:.1f} MB"
+                )
+            if epoch_skipped_steps > 0:
+                logger.warning(
+                    f"Skipped {epoch_skipped_steps} optimizer step(s) this "
+                    f"epoch due to non-finite gradients (total so far: "
+                    f"{total_skipped_steps})."
+                )
             if self.exp_path is not None:
                 self.save_ckpt(os.path.join(self.exp_path, "last.ckpt"))
 
+            epoch_record: Dict = {
+                "epoch": epoch,
+                "global_step": global_step,
+                "train_loss": sum(training_loss) / len(training_loss),
+                "epoch_time_s": round(epoch_time, 3),
+                "skipped_steps": epoch_skipped_steps,
+                **{f"train_{k}": v for k, v in vram.items()},
+            }
+
             # validation
             if val_dataloader is not None:
                 scores = self.evaluate(val_dataloader)
                 logger.info(f"--- Eval epoch-{epoch}, step-{global_step} ---")
                 for key in scores.keys():
                     logger.info("{}: {:.4f}".format(key, scores[key]))
+                epoch_record.update({f"val_{k}": v for k, v in scores.items()})
                 # save best model
                 if monitor is not None:
                     score = scores[monitor]
@@ -247,8 +357,21 @@ def train(
                             logger.info(
                                 f"Early stopping at epoch-{epoch}, step-{global_step}"
                             )
+                            metrics_history.append(epoch_record)
                             break
 
+            metrics_history.append(epoch_record)
+
+        total_time = time.perf_counter() - train_start
+        logger.info(f"--- Training complete: {total_time:.2f}s total ---")
+
+        # persist metrics history
+        if self.exp_path is not None:
+            history_path = os.path.join(self.exp_path, "metrics_history.json")
+            with open(history_path, "w") as f:
+                json.dump(metrics_history, f, indent=2)
+            logger.info(f"Metrics history saved to {history_path}")
+
         # load best model
         if load_best_model_at_last and self.exp_path is not None and os.path.isfile(
             os.path.join(self.exp_path, "best.ckpt")):
@@ -262,7 +385,7 @@ def train(
             for key in scores.keys():
                 logger.info("{}: {:.4f}".format(key, scores[key]))
 
-        return
+        return metrics_history
 
     def inference(self, dataloader, additional_outputs=None,
                   return_patient_ids=False) -> Dict[str, float]:
diff --git a/pyhealth/utils.py b/pyhealth/utils.py
index b4af8980a..efbee3977 100644
--- a/pyhealth/utils.py
+++ b/pyhealth/utils.py
@@ -21,8 +21,7 @@ def set_seed(seed):
 
 
 def create_directory(directory):
-    if not os.path.exists(directory):
-        os.makedirs(directory)
+    os.makedirs(directory, exist_ok=True)
 
 
 def load_pickle(filename):
diff --git a/pyhealth2_environment.yml b/pyhealth2_environment.yml
new file mode 100644
index 000000000..b9920a0c9
--- /dev/null
+++ b/pyhealth2_environment.yml
@@ -0,0 +1,290 @@
+name: pyhealth2
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - _python_abi3_support=1.0=hd8ed1ab_2
+  - annotated-types=0.7.0=pyhd8ed1ab_1
+  - anyio=4.14.0=pyhcf101f3_0
+  - argon2-cffi=25.1.0=pyhd8ed1ab_0
+  - argon2-cffi-bindings=25.1.0=py312h4c3975b_2
+  - arrow=1.4.0=pyhcf101f3_0
+  - asttokens=3.0.1=pyhd8ed1ab_0
+  - async-lru=2.3.0=pyhcf101f3_0
+  - attrs=26.1.0=pyhcf101f3_0
+  - babel=2.18.0=pyhcf101f3_1
+  - backports.zstd=1.6.0=py312h90b7ffd_0
+  - beautifulsoup4=4.15.0=pyha770c72_0
+  - bleach=6.4.0=pyhcf101f3_0
+  - bleach-with-css=6.4.0=hac0b51c_0
+  - brotli-python=1.2.0=py312hdb49522_1
+  - bzip2=1.0.8=h5eee18b_6
+  - c-ares=1.34.8=hb03c661_0
+  - ca-certificates=2026.6.17=hbd8a1cb_0
+  - cached-property=1.5.2=hd8ed1ab_1
+  - cached_property=1.5.2=pyha770c72_1
+  - cffi=1.17.1=py312h06ac9bb_0
+  - charset-normalizer=3.4.7=pyhd8ed1ab_0
+  - comm=0.2.3=pyhe01879c_0
+  - cpython=3.12.13=py312hd8ed1ab_0
+  - debugpy=1.8.21=py312h8285ef7_0
+  - defusedxml=0.7.1=pyhd8ed1ab_0
+  - exceptiongroup=1.3.1=pyhd8ed1ab_0
+  - executing=2.2.1=pyhd8ed1ab_0
+  - fqdn=1.5.1=pyhd8ed1ab_1
+  - gitdb=4.0.12=pyhd8ed1ab_0
+  - gitpython=3.1.50=pyhd8ed1ab_0
+  - h11=0.16.0=pyhcf101f3_1
+  - h2=4.3.0=pyhcf101f3_0
+  - hpack=4.1.0=pyhd8ed1ab_0
+  - htcondor=25.11.0=py312h7900ff3_2
+  - htcondor-classads=25.11.0=h793e66c_2
+  - htcondor-cli=25.11.0=py312h7900ff3_2
+  - htcondor-utils=25.11.0=h8d23d0f_2
+  - httpcore=1.0.9=pyh29332c3_0
+  - httpx=0.28.1=pyhd8ed1ab_0
+  - hyperframe=6.1.0=pyhd8ed1ab_0
+  - icu=78.3=h33c6efd_0
+  - importlib-metadata=9.0.0=pyhcf101f3_0
+  - ipykernel=7.3.0=pyha191276_0
+  - ipython=9.14.1=pyh53cf698_0
+  - ipython_pygments_lexers=1.1.1=pyhd8ed1ab_0
+  - isoduration=20.11.0=pyhd8ed1ab_1
+  - jedi=0.19.2=pyhd8ed1ab_1
+  - jinja2=3.1.6=pyhcf101f3_1
+  - json5=0.15.0=pyhd8ed1ab_0
+  - jsonpointer=3.1.1=pyhcf101f3_0
+  - jsonschema=4.26.0=pyhcf101f3_0
+  - jsonschema-specifications=2025.9.1=pyhcf101f3_0
+  - jsonschema-with-format-nongpl=4.26.0=hcf101f3_0
+  - jupyter-builder=1.0.2=pyhcf101f3_0
+  - jupyter-lsp=2.3.1=pyhcf101f3_0
+  - jupyter_client=8.9.1=pyhcf101f3_0
+  - jupyter_core=5.9.1=pyhc90fa1f_0
+  - jupyter_events=0.12.1=pyhcf101f3_0
+  - jupyter_server=2.20.0=pyhcf101f3_0
+  - jupyter_server_terminals=0.5.4=pyhcf101f3_0
+  - jupyterlab=4.6.0=pyhd8ed1ab_0
+  - jupyterlab_pygments=0.3.0=pyhd8ed1ab_2
+  - jupyterlab_server=2.28.0=pyhcf101f3_0
+  - keyutils=1.6.3=hb9d3cd8_0
+  - krb5=1.22.2=hbde042b_1
+  - lark=1.3.1=pyhd8ed1ab_0
+  - ld_impl_linux-64=2.44=h9e0c5a2_3
+  - libabseil=20260107.1=cxx17_h7b12aa8_0
+  - libcap=2.77=hd0affe5_1
+  - libcondor_utils=25.11.0=h4effbbe_2
+  - libcurl=8.21.0=hae6b9f4_2
+  - libdrm=2.4.125=hb03c661_1
+  - libedit=3.1.20250104=pl5321h7949ede_0
+  - libev=4.33=hd590300_2
+  - libexpat=2.7.5=h7354ed3_0
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc=15.2.0=h69a1729_7
+  - libgcc-ng=15.2.0=h166f726_7
+  - libgomp=15.2.0=h4751f2c_7
+  - liblzma=5.8.2=hb03c661_0
+  - libnghttp2=1.68.1=h877daf1_0
+  - libnsl=2.0.1=hb9d3cd8_1
+  - libpciaccess=0.18=hb9d3cd8_0
+  - libpsl=0.22.0=h49b2146_1
+  - libsodium=1.0.22=h280c20c_1
+  - libsqlite=3.53.3=h0c1763c_0
+  - libssh2=1.11.1=hcf80075_0
+  - libstdcxx=15.2.0=h39759b7_7
+  - libstdcxx-ng=15.2.0=hc03a8fd_7
+  - libsystemd0=257.13=hd0affe5_0
+  - libuuid=2.42.2=h5347b49_0
+  - libxcb=1.17.0=h9b100fa_0
+  - libxcrypt=4.4.36=hd590300_1
+  - libzlib=1.3.2=h25fd6f3_2
+  - markupsafe=3.0.3=py312h8a5da7c_1
+  - matplotlib-inline=0.2.2=pyhd8ed1ab_0
+  - mistune=3.3.2=pyhcf101f3_0
+  - munge=0.5.16=h63a00c3_0
+  - nbclient=0.11.0=pyhd8ed1ab_0
+  - nbconvert-core=7.17.1=pyhcf101f3_0
+  - nbformat=5.10.4=pyhd8ed1ab_1
+  - ncurses=6.5=h7934f7d_0
+  - nest-asyncio2=1.7.2=pyhcf101f3_0
+  - notebook=7.6.0=pyhcf101f3_0
+  - notebook-shim=0.2.4=pyhd8ed1ab_1
+  - nvidia-ml-py=13.590.48=pyhd8ed1ab_0
+  - nvitop=1.7.0=pyh707e725_0
+  - nvtop=3.3.2=h5433ab2_0
+  - openssl=3.6.3=h35e630c_0
+  - overrides=7.7.0=pyhd8ed1ab_1
+  - packaging=26.0=py312h06a4308_0
+  - pandocfilters=1.5.0=pyhd8ed1ab_0
+  - parso=0.8.7=pyhcf101f3_0
+  - pcre2=10.47=haa7fec5_0
+  - pexpect=4.9.0=pyhd8ed1ab_1
+  - pip=26.0.1=pyhc872135_1
+  - prometheus_client=0.25.0=pyhd8ed1ab_0
+  - prompt-toolkit=3.0.52=pyha770c72_0
+  - protobuf=6.33.5=py312ha7b3241_2
+  - psutil=7.2.2=py312h5253ce2_0
+  - pthread-stubs=0.3=h0ce48e5_1
+  - ptyprocess=0.7.0=pyhd8ed1ab_1
+  - pure_eval=0.2.3=pyhd8ed1ab_1
+  - pycparser=3.0=pyhcf101f3_0
+  - pygments=2.20.0=pyhd8ed1ab_0
+  - pysocks=1.7.1=pyha55dd90_7
+  - python=3.12.9=h9e4cc4f_1_cpython
+  - python-dateutil=2.9.0.post0=pyhe01879c_2
+  - python-fastjsonschema=2.21.2=pyhe01879c_0
+  - python-gil=3.12.13=hd8ed1ab_0
+  - python-htcondor=25.11.0=py312h40fa4ac_2
+  - python-json-logger=4.1.0=pyhd8ed1ab_0
+  - python-tzdata=2026.2=pyhd8ed1ab_0
+  - python_abi=3.12=3_cp312
+  - pyyaml=6.0.3=py312h8a5da7c_1
+  - pyzmq=27.1.0=py312hda471dd_3
+  - readline=8.3=hc2a1206_0
+  - referencing=0.37.0=pyhcf101f3_0
+  - rfc3339-validator=0.1.4=pyhd8ed1ab_1
+  - rfc3986-validator=0.1.1=pyh9f0ad1d_0
+  - rfc3987-syntax=1.1.0=pyhe01879c_1
+  - rpds-py=2026.5.1=py312h192e038_0
+  - scitokens-cpp=1.4.0=h096d96b_0
+  - send2trash=2.1.0=pyha191276_1
+  - sentry-sdk=2.64.0=pyhd8ed1ab_0
+  - setuptools=82.0.1=py312h06a4308_0
+  - six=1.17.0=pyhe01879c_1
+  - smmap=5.0.3=pyhcf101f3_1
+  - sniffio=1.3.1=pyhd8ed1ab_2
+  - soupsieve=2.8.4=pyhd8ed1ab_0
+  - sqlite=3.51.2=h3e8d24a_0
+  - stack_data=0.6.3=pyhd8ed1ab_1
+  - terminado=0.18.1=pyhc90fa1f_1
+  - tinycss2=1.4.0=pyhd8ed1ab_0
+  - tk=8.6.15=h54e0aa7_0
+  - tomli=2.4.1=pyhcf101f3_0
+  - traitlets=5.15.1=pyhcf101f3_0
+  - typing-extensions=4.15.0=h396c80c_0
+  - typing-inspection=0.4.2=pyhcf101f3_2
+  - typing_extensions=4.15.0=pyhcf101f3_0
+  - typing_utils=0.1.0=pyhd8ed1ab_1
+  - uri-template=1.3.0=pyhd8ed1ab_1
+  - wandb=0.28.0=py312h868fb18_0
+  - wcwidth=0.8.1=pyhd8ed1ab_0
+  - webcolors=25.10.0=pyhd8ed1ab_0
+  - webencodings=0.5.1=pyhd8ed1ab_3
+  - websocket-client=1.9.0=pyhd8ed1ab_0
+  - wheel=0.46.3=py312h06a4308_0
+  - xorg-libx11=1.8.12=h9b100fa_1
+  - xorg-libxau=1.0.12=h9b100fa_0
+  - xorg-libxdmcp=1.1.5=h9b100fa_0
+  - xorg-xorgproto=2024.1=h5eee18b_1
+  - xz=5.8.2=h448239c_0
+  - yaml=0.2.5=h280c20c_3
+  - zeromq=4.3.5=h09e67af_11
+  - zipp=4.1.0=pyhcf101f3_0
+  - zlib=1.3.2=h25fd6f3_2
+  - zstd=1.5.7=hb78ec9c_6
+  - pip:
+      - accelerate==1.13.0
+      - axial-positional-embedding==0.3.12
+      - bokeh==3.9.0
+      - boto3==1.42.88
+      - botocore==1.42.88
+      - certifi==2026.2.25
+      - click==8.3.2
+      - cloudpickle==3.1.2
+      - colt5-attention==0.11.1
+      - contourpy==1.3.3
+      - cycler==0.12.1
+      - dask==2025.11.0
+      - decorator==5.2.1
+      - distributed==2025.11.0
+      - einops==0.8.2
+      - filelock==3.25.2
+      - fonttools==4.62.1
+      - fsspec==2026.2.0
+      - hf-xet==1.4.3
+      - huggingface-hub==0.36.2
+      - hyper-connections==0.4.9
+      - idna==3.11
+      - ipywidgets==8.1.8
+      - jmespath==1.1.0
+      - joblib==1.5.3
+      - jupyterlab-widgets==3.0.16
+      - kiwisolver==1.5.0
+      - lazy-loader==0.5
+      - lightning-utilities==0.15.3
+      - linear-attention-transformer==0.19.1
+      - linformer==0.2.3
+      - litdata==0.2.61
+      - littleutils==0.2.4
+      - local-attention==1.11.2
+      - locket==1.0.0
+      - lz4==4.4.5
+      - matplotlib==3.10.8
+      - mne==1.10.2
+      - more-itertools==10.8.0
+      - mpmath==1.3.0
+      - msgpack==1.1.2
+      - narwhals==2.13.0
+      - networkx==3.6.1
+      - numpy==2.2.6
+      - nvidia-cublas-cu12==12.6.4.1
+      - nvidia-cuda-cupti-cu12==12.6.80
+      - nvidia-cuda-nvrtc-cu12==12.6.77
+      - nvidia-cuda-runtime-cu12==12.6.77
+      - nvidia-cudnn-cu12==9.5.1.17
+      - nvidia-cufft-cu12==11.3.0.4
+      - nvidia-cufile-cu12==1.11.1.6
+      - nvidia-curand-cu12==10.3.7.77
+      - nvidia-cusolver-cu12==11.7.1.2
+      - nvidia-cusparse-cu12==12.5.4.2
+      - nvidia-cusparselt-cu12==0.6.3
+      - nvidia-nccl-cu12==2.26.2
+      - nvidia-nvjitlink-cu12==12.6.85
+      - nvidia-nvtx-cu12==12.6.77
+      - obstore==0.9.2
+      - ogb==1.3.6
+      - outdated==0.2.2
+      - pandas==2.3.3
+      - partd==1.4.2
+      - peft==0.18.1
+      - pillow==12.1.1
+      - platformdirs==4.9.6
+      - polars==1.35.2
+      - polars-runtime-32==1.35.2
+      - pooch==1.9.0
+      - product-key-memory==0.3.0
+      - pyarrow==22.0.0
+      - pydantic==2.11.10
+      - pydantic-core==2.33.2
+      - pyhealth==2.0.0
+      - pyparsing==3.3.2
+      - pytz==2026.1.post1
+      - rdkit==2026.3.1
+      - regex==2026.4.4
+      - requests==2.33.1
+      - s3transfer==0.16.0
+      - safetensors==0.7.0
+      - scikit-learn==1.7.2
+      - scipy==1.17.1
+      - sortedcontainers==2.4.0
+      - sympy==1.14.0
+      - tblib==3.2.2
+      - threadpoolctl==3.6.0
+      - tifffile==2026.3.3
+      - tokenizers==0.21.4
+      - toolz==1.1.0
+      - torch==2.7.1
+      - torch-einops-utils==0.0.30
+      - torchvision==0.22.1
+      - tornado==6.5.5
+      - tqdm==4.67.3
+      - transformers==4.53.3
+      - triton==3.3.1
+      - tzdata==2026.1
+      - urllib3==2.5.0
+      - widgetsnbextension==4.0.15
+      - xyzservices==2026.3.0
+      - zict==3.0.0
+prefix: /home/wp14/miniconda3/envs/pyhealth2
diff --git a/run_table2_cc.sh b/run_table2_cc.sh
new file mode 100755
index 000000000..f313101d5
--- /dev/null
+++ b/run_table2_cc.sh
@@ -0,0 +1,104 @@
+#!/usr/bin/env bash
+# Full e2e: clone repo on CC if needed, set up conda env, sync scripts,
+# warm cache, and submit 18 Table 2 training jobs.
+#
+# Usage:
+#   bash run_table2_cc.sh           # normal run
+#   SETUP=1 bash run_table2_cc.sh   # force re-run setup even if env exists
+#   SETUP=clean bash run_table2_cc.sh  # nuke and rebuild conda env
+set -euo pipefail
+
+CC="${CC:-rianatri@cc-login.campuscluster.illinois.edu}"
+REMOTE_REPO="${REMOTE_REPO:-/u/rianatri/PyHealth}"
+LOCAL_REPO="${LOCAL_REPO:-/Users/saurabhatri/Dev/Multimodal-PyHealth}"
+SETUP="${SETUP:-auto}"  # auto | 1 | clean
+
+cd "${LOCAL_REPO}"
+
+echo "[1/5] Validate local files..."
+test -f scripts/slurm/run_table2.sh
+test -f scripts/slurm/run_cachewarm.sh
+test -f scripts/slurm/submit_table2_random.sh
+test -f scripts/slurm/setup_cc.sh
+bash -n scripts/slurm/run_table2.sh
+echo "  OK"
+
+echo "[2/5] Sync repo files to Campus Cluster..."
+# Full pyhealth source sync — ensures CC has latest model/processor/task code
+rsync -avz --relative \
+  pyhealth/ \
+  examples/mortality_prediction/unified_embedding_e2e_mimic4.py \
+  scripts/condor/warm_table2_cache.py \
+  scripts/slurm/run_table2.sh \
+  scripts/slurm/run_cachewarm.sh \
+  scripts/slurm/submit_table2_random.sh \
+  scripts/slurm/setup_cc.sh \
+  "${CC}:${REMOTE_REPO}/"
+
+echo "[3/5] Setup conda env on CC (if needed)..."
+ssh "${CC}" "REMOTE_REPO='${REMOTE_REPO}' SETUP='${SETUP}' bash -s" <<'EOF'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+
+module load miniconda3/24.9.2
+eval "$(conda shell.bash hook)"
+
+CONDA_ENV="pyhealth2"
+ENV_EXISTS=0
+conda env list | grep -q "^${CONDA_ENV}[[:space:]]" && ENV_EXISTS=1
+
+if [[ "${SETUP}" == "clean" ]]; then
+    echo "  Clean rebuild requested..."
+    bash scripts/slurm/setup_cc.sh clean
+elif [[ "${SETUP}" == "1" || "${ENV_EXISTS}" == "0" ]]; then
+    echo "  Running setup (env missing or SETUP=1)..."
+    bash scripts/slurm/setup_cc.sh
+else
+    echo "  Conda env ${CONDA_ENV} exists, skipping setup."
+    echo "  (Run with SETUP=1 to force, SETUP=clean to rebuild.)"
+    # Quick smoke test to make sure it's not broken
+    conda activate "${CONDA_ENV}"
+    python -c "import pyhealth, torch; print(f'  pyhealth OK, torch {torch.__version__}')"
+fi
+EOF
+
+echo "[4/5] Prepare remote directories + permissions..."
+ssh "${CC}" "REMOTE_REPO='${REMOTE_REPO}' bash -s" <<'EOF'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+mkdir -p logs/slurm
+chmod +x scripts/slurm/run_table2.sh
+chmod +x scripts/slurm/run_cachewarm.sh
+chmod +x scripts/slurm/submit_table2_random.sh
+chmod +x scripts/slurm/setup_cc.sh
+echo "  Permissions OK"
+EOF
+
+echo "[5/5] Submit cachewarm + chain 18 training jobs..."
+ssh "${CC}" "REMOTE_REPO='${REMOTE_REPO}' bash -s" <<'EOF'
+set -euo pipefail
+cd "${REMOTE_REPO}"
+
+CACHE_INDEX="/u/${USER}/pyhealth_cache/662e765e-a310-5499-92b3-de524ea984bb/tasks/ClinicalNotesICDLabsMIMIC4_04c5dd00-6fb7-5542-9b0d-c722773fcc42/samples_cdbbc602-34e2-5a41-8643-4c76b08829f6.ld/index.json"
+
+if [[ -f "${CACHE_INDEX}" ]]; then
+    echo "  Cache already warm — submitting 18 training jobs directly."
+    bash scripts/slurm/submit_table2_random.sh
+else
+    echo "  Cache cold — submitting cachewarm job first."
+    WARM_JOB=$(sbatch scripts/slurm/run_cachewarm.sh | awk '{print $NF}')
+    echo "  Cachewarm job ID: ${WARM_JOB}"
+    bash scripts/slurm/submit_table2_random.sh "${WARM_JOB}"
+    echo ""
+    echo "  All 19 jobs queued (1 cachewarm + 18 training)."
+    echo "  Training jobs will start automatically after cachewarm finishes."
+fi
+
+echo ""
+squeue -u rianatri --format="%.10i %.20j %.8T %.10M %.6D %R" 2>/dev/null || squeue -u rianatri
+EOF
+
+echo ""
+echo "Done. Monitor with:"
+echo "  ssh ${CC} 'squeue -u rianatri'"
+echo "  ssh ${CC} 'tail -f ${REMOTE_REPO}/logs/slurm/table2_<model>_seed<N>_<jobid>.out'"
diff --git a/scripts/compute_token_stats.py b/scripts/compute_token_stats.py
new file mode 100644
index 000000000..9e1cb0c84
--- /dev/null
+++ b/scripts/compute_token_stats.py
@@ -0,0 +1,182 @@
+#!/usr/bin/env python3
+"""Compute per-modality token counts and missing-token rates for Table 2.
+
+Iterates patients directly via dataset.iter_patients() and calls the task
+function on each one — never materializes all samples in memory, no litdata
+write, no GPU needed.
+
+Usage (on CC):
+    python scripts/compute_token_stats.py \
+        --ehr-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2 \
+        --note-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note \
+        --cache-dir /u/rianatri/pyhealth_cache \
+        --task notes_labs
+"""
+import argparse
+import os
+
+import numpy as np
+from tqdm import tqdm
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ehr-root", required=True)
+    p.add_argument("--note-root", required=True)
+    p.add_argument("--cache-dir", default=None)
+    p.add_argument("--dev", action="store_true", help="Limit to 1000 patients")
+    p.add_argument(
+        "--task",
+        type=str,
+        choices=["notes_labs"],
+        default="notes_labs",
+        help="Task to profile. 'notes_labs' uses admission-context text sections.",
+    )
+    p.add_argument(
+        "--icd-codes",
+        action="store_true",
+        default=False,
+        help="Include discharge-coded ICD codes when using --task notes_labs.",
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    os.environ.setdefault("PYHEALTH_DISABLE_DASK_DISTRIBUTED", "1")
+
+    from pyhealth.datasets import MIMIC4Dataset
+    from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+    print("Building dataset (uses cache if available)...")
+
+    ehr_tables = (
+        ["diagnoses_icd", "procedures_icd", "labevents"]
+        if args.icd_codes
+        else ["labevents"]
+    )
+    note_tables = ["discharge"]
+    task = NotesLabsMIMIC4(window_hours=24, include_icd=args.icd_codes)
+
+    kwargs = dict(
+        ehr_root=args.ehr_root,
+        note_root=args.note_root,
+        ehr_tables=ehr_tables,
+        note_tables=note_tables,
+        dev=args.dev,
+    )
+    if args.cache_dir:
+        kwargs["cache_dir"] = args.cache_dir
+
+    dataset = MIMIC4Dataset(**kwargs)
+
+    MISSING_TEXT = ""
+    LAB_CATEGORIES = task.LAB_CATEGORY_NAMES  # 10 names
+
+    # Counters
+    n_patients = n_samples = 0
+    note_total = note_missing = 0
+    note_empty_extracted = 0  # notes present but section extraction returned nothing
+    icd_total_visits = icd_missing_visits = 0
+    icd_total_codes = 0
+    lab_total_timesteps = lab_missing_timesteps = 0
+    lab_per_cat_missing = np.zeros(len(LAB_CATEGORIES), dtype=np.int64)
+
+    print("Iterating patients and accumulating stats (no litdata write)...")
+    for patient in tqdm(dataset.iter_patients(), total=len(dataset.unique_patient_ids)):
+        n_patients += 1
+        samples = task(patient)
+        if not samples:
+            continue
+        for s in samples:
+            n_samples += 1
+
+            # ── notes ────────────────────────────────────────────
+            note_texts, _ = s["admission_note_times"]
+            note_total += len(note_texts)
+            for t in note_texts:
+                if t == MISSING_TEXT:
+                    note_missing += 1
+                elif len(t) <= 1024 and t == t[:1024]:
+                    # Heuristic: if the note is exactly 1024 chars, it likely
+                    # came from the fallback raw-note path (no sections found).
+                    note_empty_extracted += 1
+
+            # ── ICD codes ────────────────────────────────────────
+            if "icd_codes" in s:
+                _, icd_visits = s["icd_codes"]
+                icd_total_visits += len(icd_visits)
+                for visit_codes in icd_visits:
+                    if visit_codes == [MISSING_TEXT]:
+                        icd_missing_visits += 1
+                    else:
+                        icd_total_codes += len(visit_codes)
+
+            # ── labs ─────────────────────────────────────────────
+            _, lab_masks = s["labs_mask"]
+            for mask_row in lab_masks:
+                lab_total_timesteps += 1
+                if not any(mask_row):
+                    lab_missing_timesteps += 1
+                for i, observed in enumerate(mask_row):
+                    if not observed:
+                        lab_per_cat_missing[i] += 1
+
+    def pct(n, d):
+        return 100.0 * n / d if d else float("nan")
+
+    print("\n" + "=" * 60)
+    print(f"TOKEN STATS — {task.task_name}")
+    print("=" * 60)
+    print(f"  Patients processed : {n_patients:>8,}")
+    print(f"  Samples (patients) : {n_samples:>8,}")
+
+    print("\n── Admission-context notes ─────────────────────────────────")
+    print(f"  Total note tokens      : {note_total:>8,}")
+    print(f"  Missing (empty str)    : {note_missing:>8,}  ({pct(note_missing, note_total):.1f}%)")
+    print(
+        f"  Fallback (raw prefix)  : {note_empty_extracted:>8,}  "
+        f"({pct(note_empty_extracted, note_total):.1f}%)"
+    )
+    print(
+        f"  Effective coverage     : {note_total - note_missing:>8,}  "
+        f"({pct(note_total - note_missing, note_total):.1f}%)"
+    )
+
+    if "icd_codes" in s:
+        print("\n── ICD Codes ────────────────────────────────────────────────")
+        print(f"  Total visit tokens  : {icd_total_visits:>8,}")
+        print(
+            f'  Missing visits [""] : {icd_missing_visits:>8,}  '
+            f"({pct(icd_missing_visits, icd_total_visits):.1f}%)"
+        )
+        print(f"  Total ICD codes     : {icd_total_codes:>8,}  (in non-missing visits)")
+
+    print("\n── Labs ─────────────────────────────────────────────────────")
+    print(f"  Total timestep tokens : {lab_total_timesteps:>8,}")
+    print(
+        f"  Missing timesteps     : {lab_missing_timesteps:>8,}  "
+        f"({pct(lab_missing_timesteps, lab_total_timesteps):.1f}%)"
+    )
+    print(f"\n  Per-category missingness (across all timesteps):")
+    for i, cat in enumerate(LAB_CATEGORIES):
+        n_miss = int(lab_per_cat_missing[i])
+        print(
+            f"    {cat:<15}: {n_miss:>8,} / {lab_total_timesteps:>8,}  "
+            f"({pct(n_miss, lab_total_timesteps):.1f}% missing)"
+        )
+
+    print("\n── Summary ──────────────────────────────────────────────────")
+    total_tokens = note_total + icd_total_visits + lab_total_timesteps
+    total_missing = note_missing + icd_missing_visits + lab_missing_timesteps
+    print(f"  All modalities total tokens  : {total_tokens:>8,}")
+    print(
+        f"  All modalities missing tokens: {total_missing:>8,}  "
+        f"({pct(total_missing, total_tokens):.1f}%)"
+    )
+    print()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/condor/batch_ablation.sub b/scripts/condor/batch_ablation.sub
new file mode 100644
index 000000000..04c650b64
--- /dev/null
+++ b/scripts/condor/batch_ablation.sub
@@ -0,0 +1,100 @@
+# HTCondor submission — PyHealth Ablation Runs
+#
+# Four conditions on devsplit (1000 patients), MLP, 100 epochs:
+#   1. baseline:  notes+labs only (ICD off, vitals off, balanced off)
+#   2. icd_on:    notes+labs+ICD codes (leakage ablation)
+#   3. vitals_on: notes+labs+vitals (chartevents ablation)
+#   4. balanced:  notes+labs with balanced sampling (1:1 pos:neg)
+#
+# All use MLP backbone, seed=42, embedding_dim=128, patience=10.
+# Jobs start with "PyHealth" for easy queue filtering.
+#
+# To submit from the project root on the submit host:
+#   cd /home/rianatri/PyHealth && condor_submit scripts/condor/batch_ablation.sub
+
+initialdir           = /home/rianatri/PyHealth
+executable           = /home/rianatri/PyHealth/scripts/condor/run_table2.sh
+transfer_executable  = False
+arguments            = $(model) $(seed)
+getenv               = True
+
++JobBatchName = "PyHealth"
+
+output = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).out
+error  = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).err
+log    = logs/condor/ablation_$(ClusterId)_$(Process)_$(model)_seed$(seed).log
+
+request_gpus   = 1
+request_cpus   = 4
+request_memory = 48000
+
+Rank = TARGET.CUDAGlobalMemoryMb
+
+# ── Condition 1: baseline (notes+labs, no ICD, no vitals, no balancing) ────────
+queue model, seed, environment from (
+    mlp, 42, "TABLE2_TASK=notes_labs \
+               TABLE2_RUN_LABEL=PyHealth_baseline \
+               TABLE2_DEV_MODE=1 \
+               TABLE2_DEV_COUNT=1000 \
+               TABLE2_EPOCHS=100 \
+               TABLE2_NUM_WORKERS=2 \
+               TABLE2_FREEZE_ENCODER=0 \
+               TABLE2_ICD_CODES=0 \
+               TABLE2_INCLUDE_VITALS=0 \
+               TABLE2_BALANCED_SAMPLING=0 \
+               TABLE2_OUTPUT_DIR=output/ablation_baseline \
+               TABLE2_WINDOW_HOURS=24 \
+               TABLE2_PATIENCE=10"
+)
+
+# ── Condition 2: icd_on (notes+labs+ICD codes — discharge-coded leakage) ──────
+queue model, seed, environment from (
+    mlp, 42, "TABLE2_TASK=notes_labs \
+               TABLE2_RUN_LABEL=PyHealth_icd_on \
+               TABLE2_DEV_MODE=1 \
+               TABLE2_DEV_COUNT=1000 \
+               TABLE2_EPOCHS=100 \
+               TABLE2_NUM_WORKERS=2 \
+               TABLE2_FREEZE_ENCODER=0 \
+               TABLE2_ICD_CODES=1 \
+               TABLE2_INCLUDE_VITALS=0 \
+               TABLE2_BALANCED_SAMPLING=0 \
+               TABLE2_OUTPUT_DIR=output/ablation_icd_on \
+               TABLE2_WINDOW_HOURS=24 \
+               TABLE2_PATIENCE=10"
+)
+
+# ── Condition 3: vitals_on (notes+labs+vitals from chartevents) ───────────────
+queue model, seed, environment from (
+    mlp, 42, "TABLE2_TASK=notes_labs \
+               TABLE2_RUN_LABEL=PyHealth_vitals_on \
+               TABLE2_DEV_MODE=1 \
+               TABLE2_DEV_COUNT=1000 \
+               TABLE2_EPOCHS=100 \
+               TABLE2_NUM_WORKERS=2 \
+               TABLE2_FREEZE_ENCODER=0 \
+               TABLE2_ICD_CODES=0 \
+               TABLE2_INCLUDE_VITALS=1 \
+               TABLE2_BALANCED_SAMPLING=0 \
+               TABLE2_OUTPUT_DIR=output/ablation_vitals_on \
+               TABLE2_WINDOW_HOURS=24 \
+               TABLE2_PATIENCE=10"
+)
+
+# ── Condition 4: balanced (notes+labs with 1:1 pos:neg training sampling) ─────
+queue model, seed, environment from (
+    mlp, 42, "TABLE2_TASK=notes_labs \
+               TABLE2_RUN_LABEL=PyHealth_balanced \
+               TABLE2_DEV_MODE=1 \
+               TABLE2_DEV_COUNT=1000 \
+               TABLE2_EPOCHS=100 \
+               TABLE2_NUM_WORKERS=2 \
+               TABLE2_FREEZE_ENCODER=0 \
+               TABLE2_ICD_CODES=0 \
+               TABLE2_INCLUDE_VITALS=0 \
+               TABLE2_BALANCED_SAMPLING=1 \
+               TABLE2_BALANCED_RATIO=1.0 \
+               TABLE2_OUTPUT_DIR=output/ablation_balanced \
+               TABLE2_WINDOW_HOURS=24 \
+               TABLE2_PATIENCE=10"
+)
\ No newline at end of file
diff --git a/scripts/condor/batch_noteslabs_poc.sub b/scripts/condor/batch_noteslabs_poc.sub
new file mode 100644
index 000000000..8ae75df70
--- /dev/null
+++ b/scripts/condor/batch_noteslabs_poc.sub
@@ -0,0 +1,45 @@
+# HTCondor submission — notes_labs POC (single seed, lightweight params)
+#
+# Purpose: proof-of-concept run for the NotesLabsMIMIC4 task (admission-context
+# note sections + StageNet labs, no ICD leakage).  Single seed; no multi-seed
+# coverage needed for a POC.  Models chosen for low VRAM: mlp, rnn,
+# bottleneck_transformer.  BERT encoder is frozen (--freeze-encoder) to halve
+# BERT VRAM and keep each job well under 24 GB (A10 / A40 compatible).
+#
+# Ablation variants:
+#   noteslabs_poc       — recommended: notes + labs, no ICD codes
+#   noteslabs_icd_poc   — ablation: notes + labs + discharge-coded ICD codes
+#
+# To submit:
+#   condor_submit scripts/condor/batch_noteslabs_poc.sub
+
+executable = scripts/condor/run_table2.sh
+arguments  = $(model) $(seed)
+getenv     = True
+
+# notes_labs task, frozen encoder, ablation flag controlled per queue block
+environment = "TABLE2_TASK=notes_labs \
+               TABLE2_OUTPUT_DIR=/home/rianatri/noteslabs_poc \
+               TABLE2_RUN_LABEL=noteslabs_poc \
+               TABLE2_EPOCHS=10 \
+               TABLE2_FREEZE_ENCODER=1 \
+               TABLE2_WINDOW_HOURS=24"
+
+output = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).out
+error  = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).err
+log    = logs/condor/noteslabs_poc_$(ClusterId)_$(Process)_$(model)_seed$(seed).log
+
+# 24 GB covers A10/A40 with frozen BERT + embedding_dim=64.
+# For bottleneck_transformer bump to 40 GB if needed on smaller slots.
+request_gpus   = 1
+request_cpus   = 4
+request_memory = 48000
+
+Rank = TARGET.CUDAGlobalMemoryMb
+
+# ── noteslabs_poc: notes + labs, no ICD codes ───────────────────────────────
+queue model, seed from (
+    mlp,                   42
+    rnn,                   42
+    bottleneck_transformer, 42
+)
diff --git a/scripts/condor/batch_table2.sub b/scripts/condor/batch_table2.sub
new file mode 100644
index 000000000..d044237ec
--- /dev/null
+++ b/scripts/condor/batch_table2.sub
@@ -0,0 +1,46 @@
+# HTCondor submission file — Table 2: Structured EHR + Clinical Notes
+# Benchmarks all 6 backbone architectures on the clinical_notes_icd_labs
+# task (ICD codes + lab events + discharge/radiology notes).
+#
+# Each job claims one A100 80GB on sunlab-c01 and trains the assigned
+# model for 3 seeds sequentially.  6 jobs = 6 A100s in parallel.
+#
+# Prerequisites:
+#   1. Fill in CONDA_ENV, PROJECT_DIR, EHR_ROOT, NOTE_ROOT, CACHE_DIR
+#      (and optionally CONDA_SH) inside scripts/condor/run_table2.sh
+#   2. mkdir -p logs/condor
+#
+# Submit:
+#   condor_submit scripts/condor/batch_table2.sub
+#
+# Monitor:
+#   condor_q
+#   condor_status -gpus
+
+executable = scripts/condor/run_table2.sh
+arguments  = $(model)
+getenv = True
+
+# Per-job output / error / shared Condor log
+output = logs/condor/table2_$(ClusterId)_$(model).out
+error  = logs/condor/table2_$(ClusterId)_$(model).err
+log    = logs/condor/table2_$(ClusterId).log
+
+# Resource requests
+request_gpus   = 1
+request_cpus   = 8
+request_memory = 32768MB
+request_disk   = 20GB
+
+# ── Target: sunlab-c01 (8 × NVIDIA A100 80GB PCIe) ───────────────
+Requirements = (Machine == "sunlab-c01.cs.illinois.edu")
+
+# ── Queue one job per backbone model ─────────────────────────────
+queue model from (
+    mlp
+    rnn
+    transformer
+    bottleneck_transformer
+    ehrmamba
+    jambaehr
+)
diff --git a/scripts/condor/run_table2.sh b/scripts/condor/run_table2.sh
new file mode 100755
index 000000000..09d2a5d57
--- /dev/null
+++ b/scripts/condor/run_table2.sh
@@ -0,0 +1,249 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+MODEL="${1:?usage: run_table2.sh <model> <seed>}"
+SEED="${2:?usage: run_table2.sh <model> <seed>}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/home/rianatri/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/home/rianatri/pyhealth_cache}"
+CONDA_SH="${CONDA_SH:-}"
+PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}"
+TABLE2_EPOCHS="${TABLE2_EPOCHS:-20}"
+TABLE2_NUM_WORKERS="${TABLE2_NUM_WORKERS:-2}"
+TABLE2_DEV_MODE="${TABLE2_DEV_MODE:-0}"
+TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}"
+TABLE2_WINDOW_HOURS="${TABLE2_WINDOW_HOURS:-24}"
+TABLE2_OUTPUT_DIR="${TABLE2_OUTPUT_DIR:-output/table2}"
+TABLE2_RUN_LABEL="${TABLE2_RUN_LABEL:-full}"
+TABLE2_FREEZE_ENCODER="${TABLE2_FREEZE_ENCODER:-0}"
+TABLE2_ICD_CODES="${TABLE2_ICD_CODES:-0}"
+TABLE2_INCLUDE_VITALS="${TABLE2_INCLUDE_VITALS:-0}"
+TABLE2_BALANCED_SAMPLING="${TABLE2_BALANCED_SAMPLING:-0}"
+TABLE2_BALANCED_RATIO="${TABLE2_BALANCED_RATIO:-1.0}"
+
+resolve_conda_sh() {
+    if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [[ -f "${c}" ]]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [[ -n "${found}" && -f "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then
+    echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2
+    exit 1
+fi
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+
+JOB_TAG="${MODEL}_seed${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}"
+JOB_CACHE_DIR="${CACHE_DIR}/${JOB_TAG}"
+mkdir -p "${JOB_CACHE_DIR}"
+
+if [[ -n "${_CONDOR_SCRATCH_DIR:-}" ]]; then
+    export DASK_TEMPORARY_DIRECTORY="${_CONDOR_SCRATCH_DIR}/dask-${JOB_TAG}"
+else
+    export DASK_TEMPORARY_DIRECTORY="/tmp/dask-${JOB_TAG}"
+fi
+mkdir -p "${DASK_TEMPORARY_DIRECTORY}"
+# Ensure local repo package is importable even if not installed into env.
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+export PYHEALTH_DISABLE_DASK_DISTRIBUTED
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:256}"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash setup.sh" >&2
+    exit 1
+fi
+
+echo "========================================================"
+echo "  Table 2 run  |  label=${TABLE2_RUN_LABEL}"
+echo "  Model     : ${MODEL}"
+echo "  Seed      : ${SEED}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Conda env: ${CONDA_ENV}"
+echo "  Conda sh : ${CONDA_SH}"
+echo "  EHR root : ${EHR_ROOT}"
+echo "  Note root: ${NOTE_ROOT}"
+echo "  Cache dir: ${CACHE_DIR}"
+echo "  Job cache : ${JOB_CACHE_DIR}"
+echo "  Dask temp: ${DASK_TEMPORARY_DIRECTORY}"
+echo "  Dask dist: ${PYHEALTH_DISABLE_DASK_DISTRIBUTED} (1=local scheduler)"
+echo "  Epochs   : ${TABLE2_EPOCHS}"
+echo "  Workers  : ${TABLE2_NUM_WORKERS}"
+echo "  Dev mode : ${TABLE2_DEV_MODE}"
+echo "  Task      : ${TABLE2_TASK}"
+echo "  Window    : ${TABLE2_WINDOW_HOURS}h"
+echo "  Output dir: ${TABLE2_OUTPUT_DIR}"
+echo "  Patience  : 5 (early stopping)"
+echo "  Freeze enc: ${TABLE2_FREEZE_ENCODER} (1=freeze BERT)"
+echo "  ICD codes : ${TABLE2_ICD_CODES} (1=include, ablation only)"
+echo "  Vitals    : ${TABLE2_INCLUDE_VITALS} (1=include chartevents)"
+echo "  Balanced  : ${TABLE2_BALANCED_SAMPLING} (1=undersample negatives)"
+echo "  Bal ratio : ${TABLE2_BALANCED_RATIO}"
+echo "========================================================"
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --note-root "${NOTE_ROOT}"
+    --cache-dir "${CACHE_DIR}"
+    --task "${TABLE2_TASK}"
+    --observation-window-hours "${TABLE2_WINDOW_HOURS}"
+    --model "${MODEL}"
+    --embedding-dim 128
+    --hidden-dim 128
+    --heads 4
+    --num-layers 2
+    --dropout 0.1
+    --epochs "${TABLE2_EPOCHS}"
+    --batch-size 16
+    --weight-decay 1e-5
+    --num-workers "${TABLE2_NUM_WORKERS}"
+    --output-dir "${TABLE2_OUTPUT_DIR}"
+    --patience "${TABLE2_PATIENCE:-5}"
+)
+
+if [[ "${TABLE2_DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev "${TABLE2_DEV_COUNT:-1000}")
+fi
+
+if [[ "${TABLE2_FREEZE_ENCODER}" == "1" ]]; then
+    COMMON+=(--freeze-encoder)
+fi
+
+if [[ "${TABLE2_ICD_CODES}" == "1" ]]; then
+    COMMON+=(--icd-codes)
+fi
+
+if [[ "${TABLE2_INCLUDE_VITALS}" == "1" ]]; then
+    COMMON+=(--include-vitals)
+fi
+
+if [[ "${TABLE2_BALANCED_SAMPLING}" == "1" ]]; then
+    COMMON+=(--balanced-sampling --balanced-ratio "${TABLE2_BALANCED_RATIO}")
+fi
+
+if [[ "${TABLE2_DRY_RUN:-0}" == "1" ]]; then
+    echo "Dry-run complete: conda activation and argument assembly succeeded."
+    exit 0
+fi
+
+case "${MODEL}" in
+    mlp)
+        # No dim overrides needed; A6000/A100 handles bs=16 at embedding-dim=128.
+        COMMON+=(--batch-size "${TABLE2_BS_MLP:-16}")
+        ;;
+    rnn)
+        COMMON+=(--batch-size "${TABLE2_BS_RNN:-16}")
+        ;;
+    transformer)
+        # Bumped batch-size 1→2: at bs=1 ~9000s/epoch × 20 = 50h, exceeds deadline.
+        # embedding-dim=64 is tiny; bs=2 is safe on A6000 (48GB) and A100 (80GB).
+        COMMON+=(
+            --batch-size 2
+            --embedding-dim 64
+            --hidden-dim 64
+            --heads 2
+            --num-layers 1
+        )
+        ;;
+    bottleneck_transformer)
+        # Same timing issue as transformer; bs=2 is safe given embedding-dim=96.
+        COMMON+=(
+            --batch-size 2
+            --embedding-dim 96
+            --hidden-dim 96
+            --heads 2
+            --num-layers 1
+            --max-grad-norm 0.5
+            --bottlenecks-n 4
+            --fusion-startidx 1
+        )
+        ;;
+    ehrmamba)
+        COMMON+=(
+            --batch-size 2
+            --embedding-dim 96
+            --hidden-dim 96
+            --mamba-state-size 16
+            --mamba-conv-kernel 4
+        )
+        ;;
+    jambaehr)
+        # Bumped batch-size 1→2 for same timing reason as transformer.
+        # Reduced jamba-mamba-layers 4→2: 3 total layers still a valid Jamba model,
+        # saves ~30-40% per-step cost.
+        COMMON+=(
+            --batch-size 2
+            --embedding-dim 64
+            --hidden-dim 64
+            --jamba-transformer-layers 1
+            --jamba-mamba-layers 2
+            --mamba-state-size 16
+            --mamba-conv-kernel 4
+        )
+        ;;
+esac
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" --seed "${SEED}"
+
+echo "========================================================"
+echo "  Completed label=${TABLE2_RUN_LABEL} model=${MODEL} seed=${SEED}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/condor/submit_table2_random.sh b/scripts/condor/submit_table2_random.sh
new file mode 100755
index 000000000..c9c8a547b
--- /dev/null
+++ b/scripts/condor/submit_table2_random.sh
@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+SUB_TEMPLATE="${SUB_TEMPLATE:-scripts/condor/batch_table2.sub}"
+GENERATED_SUB="${GENERATED_SUB:-scripts/condor/batch_table2.generated.sub}"
+SEED_MANIFEST="${SEED_MANIFEST:-scripts/condor/table2_random_seeds.txt}"
+
+cd "${PROJECT_DIR}"
+
+mkdir -p logs/condor
+
+mapfile -t SEEDS < <(
+python3 - <<'PY'
+import random
+
+rng = random.SystemRandom()
+for seed in rng.sample(range(1, 2_147_483_647), 3):
+    print(seed)
+PY
+)
+
+if [[ "${#SEEDS[@]}" -ne 3 ]]; then
+    echo "ERROR: failed to generate 3 random seeds." >&2
+    exit 1
+fi
+
+printf '%s\n' "${SEEDS[@]}" > "${SEED_MANIFEST}"
+
+S0="${SEEDS[0]}"
+S1="${SEEDS[1]}"
+S2="${SEEDS[2]}"
+
+# Replace SEED0/SEED1/SEED2 placeholders in the template.
+# The template has two queue blocks (heavy/light) each with their own Rank —
+# sed replaces placeholders in both blocks in one pass.
+sed \
+    -e "s/SEED0/${S0}/g" \
+    -e "s/SEED1/${S1}/g" \
+    -e "s/SEED2/${S2}/g" \
+    "${SUB_TEMPLATE}" > "${GENERATED_SUB}"
+
+echo "Generated random seeds: ${SEEDS[*]}"
+echo "Seed manifest: ${SEED_MANIFEST}"
+echo "Submit file  : ${GENERATED_SUB}"
+
+condor_submit "${GENERATED_SUB}"
diff --git a/scripts/condor/warm_table2_cache.py b/scripts/condor/warm_table2_cache.py
new file mode 100644
index 000000000..c7dc8ddb8
--- /dev/null
+++ b/scripts/condor/warm_table2_cache.py
@@ -0,0 +1,34 @@
+#!/usr/bin/env python3
+"""Warm the shared Table 2 cache from a real Python file.
+
+This exists because Python 3.12 multiprocessing with spawn cannot safely
+re-import a `python - <<'PY'` stdin script as `__main__`.
+"""
+
+from __future__ import annotations
+
+import os
+
+from pyhealth.datasets import MIMIC4Dataset
+from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+
+def main() -> None:
+    cache_root = os.environ["TABLE2_SHARED_CACHE_ROOT"]
+    dataset = MIMIC4Dataset(
+        ehr_root=os.environ["EHR_ROOT"],
+        ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"],
+        note_root=os.environ["NOTE_ROOT"],
+        note_tables=["discharge", "radiology"],
+        cache_dir=cache_root,
+        dev=os.environ.get("TABLE2_DEV_MODE", "0") == "1",
+        num_workers=int(os.environ["TABLE2_CACHE_WARM_NUM_WORKERS"]),
+    )
+    dataset.set_task(
+        NotesLabsMIMIC4(window_hours=24, include_icd=True),
+        num_workers=int(os.environ["TABLE2_CACHE_WARM_NUM_WORKERS"]),
+    )
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/slurm/run_cachewarm.sh b/scripts/slurm/run_cachewarm.sh
new file mode 100755
index 000000000..16ecbbe35
--- /dev/null
+++ b/scripts/slurm/run_cachewarm.sh
@@ -0,0 +1,56 @@
+#!/usr/bin/env bash
+#SBATCH --job-name=table2_cachewarm
+#SBATCH --account=jimeng-cs-eng
+#SBATCH --partition=eng-research-gpu
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --cpus-per-task=4
+#SBATCH --mem=48G
+#SBATCH --gres=gpu:1
+#SBATCH --time=08:00:00
+#SBATCH --output=logs/slurm/table2_cachewarm_%j.out
+#SBATCH --error=logs/slurm/table2_cachewarm_%j.err
+# Requests 1 GPU even though cachewarm is CPU-only — eng-research-gpu
+# partition requires --gres=gpu to schedule. GPU will sit idle.
+set -euo pipefail
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/u/rianatri/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/u/${USER}/pyhealth_cache}"
+NUM_WORKERS="${TABLE2_NUM_WORKERS:-4}"
+
+module load miniconda3/24.9.2
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+cd "${PROJECT_DIR}"
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}"
+# Never use /tmp for dask temp — it fills up and corrupts shared parquet cache.
+export DASK_TEMPORARY_DIRECTORY="${SLURM_TMPDIR:-/u/${USER}/dask_tmp}/dask-cachewarm-${SLURM_JOB_ID:-local}"
+mkdir -p "${DASK_TEMPORARY_DIRECTORY}"
+
+echo "========================================================"
+echo "  Table 2 cache warm"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Node      : $(hostname)"
+echo "  EHR root  : ${EHR_ROOT}"
+echo "  Note root : ${NOTE_ROOT}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  Workers   : ${NUM_WORKERS}"
+echo "========================================================"
+
+TABLE2_SHARED_CACHE_ROOT="${CACHE_DIR}" \
+EHR_ROOT="${EHR_ROOT}" \
+NOTE_ROOT="${NOTE_ROOT}" \
+TABLE2_CACHE_WARM_NUM_WORKERS="${NUM_WORKERS}" \
+TABLE2_DEV_MODE="0" \
+    python scripts/condor/warm_table2_cache.py
+# Exit 0 even if cache was already warm — ensures --dependency=afterok works.
+exit 0
+
+echo "========================================================"
+echo "  Cache warm complete"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/slurm/run_table2.sh b/scripts/slurm/run_table2.sh
new file mode 100755
index 000000000..8e7356922
--- /dev/null
+++ b/scripts/slurm/run_table2.sh
@@ -0,0 +1,171 @@
+#!/usr/bin/env bash
+# SLURM runner — Table 2, one model + one seed per job.
+# Submitted by submit_table2_random.sh via sbatch with --export=MODEL=...,SEED=...
+# Resource flags (#SBATCH) are passed on the sbatch command line, not here,
+# because heavy vs light models have different memory/time requirements.
+set -euo pipefail
+
+MODEL="${MODEL:?MODEL env var required}"
+SEED="${SEED:?SEED env var required}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/u/rianatri/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/u/${USER}/pyhealth_cache}"
+TABLE2_EPOCHS="${TABLE2_EPOCHS:-20}"
+TABLE2_NUM_WORKERS="${TABLE2_NUM_WORKERS:-2}"
+TABLE2_DEV_MODE="${TABLE2_DEV_MODE:-0}"
+TABLE2_TASK="${TABLE2_TASK:-clinical_notes_icd_labs}"
+TABLE2_WINDOW_HOURS="${TABLE2_WINDOW_HOURS:-24}"
+TABLE2_OUTPUT_DIR="${TABLE2_OUTPUT_DIR:-output/table2}"
+TABLE2_RUN_LABEL="${TABLE2_RUN_LABEL:-full}"
+TABLE2_FREEZE_ENCODER="${TABLE2_FREEZE_ENCODER:-0}"
+TABLE2_ICD_CODES="${TABLE2_ICD_CODES:-0}"
+
+# ── Activate conda ────────────────────────────────────────────────
+module load miniconda3/24.9.2
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+
+JOB_TAG="${MODEL}_seed${SEED}_j${SLURM_JOB_ID:-local}"
+
+# Use SLURM local scratch if available (--tmp= requested), else shared home.
+# Never use /tmp — it fills up and corrupts shared parquet cache.
+if [[ -n "${SLURM_TMPDIR:-}" ]]; then
+    export DASK_TEMPORARY_DIRECTORY="${SLURM_TMPDIR}/dask-${JOB_TAG}"
+else
+    export DASK_TEMPORARY_DIRECTORY="/u/${USER}/dask_tmp/dask-${JOB_TAG}"
+fi
+mkdir -p "${DASK_TEMPORARY_DIRECTORY}"
+
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}"
+export PYTORCH_CUDA_ALLOC_CONF="${PYTORCH_CUDA_ALLOC_CONF:-expandable_segments:True,max_split_size_mb:256}"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash setup.sh" >&2
+    exit 1
+fi
+
+echo "========================================================"
+echo "  Table 2 run  |  label=${TABLE2_RUN_LABEL}"
+echo "  Model     : ${MODEL}"
+echo "  Seed      : ${SEED}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  SLURM job : ${SLURM_JOB_ID:-local}"
+echo "  Node      : $(hostname)"
+echo "  Conda env : ${CONDA_ENV}"
+echo "  EHR root  : ${EHR_ROOT}"
+echo "  Note root : ${NOTE_ROOT}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  Dask temp : ${DASK_TEMPORARY_DIRECTORY}"
+echo "  Epochs    : ${TABLE2_EPOCHS}"
+echo "  Workers   : ${TABLE2_NUM_WORKERS}"
+echo "  Dev mode  : ${TABLE2_DEV_MODE}"
+echo "  Output dir: ${TABLE2_OUTPUT_DIR}"
+echo "  Task      : ${TABLE2_TASK}"
+echo "  Window    : ${TABLE2_WINDOW_HOURS}h"
+echo "  Patience  : 5 (early stopping)"
+echo "  Freeze enc: ${TABLE2_FREEZE_ENCODER} (1=freeze BERT)"
+echo "  ICD codes : ${TABLE2_ICD_CODES} (1=include, ablation only)"
+echo "========================================================"
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --note-root "${NOTE_ROOT}"
+    --cache-dir "${CACHE_DIR}"
+    --task "${TABLE2_TASK}"
+    --observation-window-hours "${TABLE2_WINDOW_HOURS}"
+    --model "${MODEL}"
+    --embedding-dim 128
+    --hidden-dim 128
+    --heads 4
+    --num-layers 2
+    --dropout 0.1
+    --epochs "${TABLE2_EPOCHS}"
+    --batch-size 16
+    --weight-decay 1e-5
+    --num-workers "${TABLE2_NUM_WORKERS}"
+    --output-dir "${TABLE2_OUTPUT_DIR}"
+    --patience 5
+)
+
+if [[ "${TABLE2_DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev)
+fi
+
+if [[ "${TABLE2_FREEZE_ENCODER}" == "1" ]]; then
+    COMMON+=(--freeze-encoder)
+fi
+
+if [[ "${TABLE2_ICD_CODES}" == "1" ]]; then
+    COMMON+=(--icd-codes)
+fi
+
+if [[ "${TABLE2_DRY_RUN:-0}" == "1" ]]; then
+    echo "Dry-run complete."
+    exit 0
+fi
+
+case "${MODEL}" in
+    mlp)
+        # BERT encoder dominates VRAM even for lightweight heads.
+        # A10 24GB: bs=4 safe. A100/H200: bs=16 default is fine.
+        COMMON+=(--batch-size "${TABLE2_BS_MLP:-4}")
+        ;;
+    rnn)
+        COMMON+=(--batch-size "${TABLE2_BS_RNN:-4}")
+        ;;
+    transformer)
+        # Default bs=1 for A10 24GB; override via TABLE2_BS_TRANSFORMER for larger GPUs.
+        COMMON+=(
+            --batch-size "${TABLE2_BS_TRANSFORMER:-1}"
+            --embedding-dim 64
+            --hidden-dim 64
+            --heads 2
+            --num-layers 1
+        )
+        ;;
+    bottleneck_transformer)
+        COMMON+=(
+            --batch-size "${TABLE2_BS_BOTTLENECK:-1}"
+            --embedding-dim 96
+            --hidden-dim 96
+            --heads 2
+            --num-layers 1
+            --max-grad-norm 0.5
+            --bottlenecks-n 4
+            --fusion-startidx 1
+        )
+        ;;
+    ehrmamba)
+        COMMON+=(
+            --batch-size "${TABLE2_BS_EHRMAMBA:-2}"
+            --embedding-dim 96
+            --hidden-dim 96
+            --mamba-state-size 16
+            --mamba-conv-kernel 4
+        )
+        ;;
+    jambaehr)
+        COMMON+=(
+            --batch-size "${TABLE2_BS_JAMBAEHR:-1}"
+            --embedding-dim 64
+            --hidden-dim 64
+            --jamba-transformer-layers 1
+            --jamba-mamba-layers 2
+            --mamba-state-size 16
+            --mamba-conv-kernel 4
+        )
+        ;;
+esac
+
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}" --seed "${SEED}"
+
+echo "========================================================"
+echo "  Completed label=${TABLE2_RUN_LABEL} model=${MODEL} seed=${SEED}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/slurm/setup_cc.sh b/scripts/slurm/setup_cc.sh
new file mode 100755
index 000000000..805c13775
--- /dev/null
+++ b/scripts/slurm/setup_cc.sh
@@ -0,0 +1,90 @@
+#!/usr/bin/env bash
+# One-time setup for PyHealth on the UIUC Campus Cluster.
+# Run this once from the login node after cloning the repo:
+#
+#   ssh rianatri@cc-login.campuscluster.illinois.edu
+#   cd ~/PyHealth
+#   bash scripts/slurm/setup_cc.sh
+#
+# Pass "clean" to nuke and rebuild the conda env from scratch:
+#   bash scripts/slurm/setup_cc.sh clean
+set -euo pipefail
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-$(cd "$(dirname "$0")/../.." && pwd)}"
+CACHE_DIR="${CACHE_DIR:-/u/${USER}/pyhealth_cache}"
+
+cd "${PROJECT_DIR}"
+
+echo "========================================================"
+echo "  PyHealth CC setup"
+echo "  Conda env : ${CONDA_ENV}"
+echo "  Project   : ${PROJECT_DIR}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
+
+# ── Load conda ────────────────────────────────────────────────────
+module load miniconda3/24.9.2
+eval "$(conda shell.bash hook)"
+
+# ── Clean if requested ────────────────────────────────────────────
+if [[ "${1:-}" == "clean" ]]; then
+    echo "[clean] Removing existing conda env ${CONDA_ENV}..."
+    conda env remove -n "${CONDA_ENV}" -y 2>/dev/null || true
+fi
+
+# ── Create env if it doesn't exist ───────────────────────────────
+if ! conda env list | grep -q "^${CONDA_ENV}[[:space:]]"; then
+    echo "[1/4] Creating conda env ${CONDA_ENV} (Python 3.12)..."
+    conda create -n "${CONDA_ENV}" python=3.12 -y
+else
+    echo "[1/4] Conda env ${CONDA_ENV} already exists, skipping create."
+fi
+
+conda activate "${CONDA_ENV}"
+
+# ── Install PyTorch with CUDA 12.x (matches A10 driver on CC) ────
+echo "[2/4] Installing PyTorch 2.7.x + CUDA 12.1..."
+pip install torch==2.7.1 torchvision --index-url https://download.pytorch.org/whl/cu121 --quiet
+
+# ── Install pyhealth and all dependencies in editable mode ───────
+echo "[3/4] Installing PyHealth (editable) + all deps..."
+pip install -e ".[full]" --quiet 2>/dev/null || pip install -e . --quiet
+
+# Install any extras not covered by pyproject.toml
+pip install platformdirs filelock --quiet
+
+# ── Pre-download Bio_ClinicalBERT so jobs don't race to HuggingFace ──
+echo "[4/4] Pre-caching Bio_ClinicalBERT tokenizer and weights..."
+python - <<'PY'
+from transformers import AutoTokenizer, AutoModel
+print("  Downloading Bio_ClinicalBERT...")
+AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+print("  Done.")
+PY
+
+# ── Create cache and log dirs ─────────────────────────────────────
+mkdir -p "${CACHE_DIR}"
+mkdir -p "${PROJECT_DIR}/logs/slurm"
+
+# ── Smoke test ────────────────────────────────────────────────────
+echo ""
+echo "Smoke test..."
+python -c "
+import torch, pyhealth, transformers
+print(f'  torch      : {torch.__version__}')
+print(f'  cuda avail : {torch.cuda.is_available()}')
+print(f'  pyhealth   : {pyhealth.__version__}')
+print(f'  transformers: {transformers.__version__}')
+import pyhealth.models.ehrmamba, pyhealth.models.jamba_ehr
+print('  EHRMamba   : OK')
+print('  JambaEHR   : OK')
+"
+
+echo ""
+echo "========================================================"
+echo "  Setup complete. Activate with:"
+echo "    module load miniconda3/24.9.2 && conda activate ${CONDA_ENV}"
+echo "========================================================"
diff --git a/scripts/slurm/submit_bottleneck_transformer.sh b/scripts/slurm/submit_bottleneck_transformer.sh
new file mode 100755
index 000000000..ed8612c69
--- /dev/null
+++ b/scripts/slurm/submit_bottleneck_transformer.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for bottleneck_transformer — 3 fixed seeds, IllinoisComputes-GPU.
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_bottleneck_transformer.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_bottleneck_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=18:00:00 \
+        --output="logs/slurm/table2_bottleneck_transformer_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_bottleneck_transformer_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=bottleneck_transformer,SEED="${seed}",TABLE2_BS_BOTTLENECK=4 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted bottleneck_transformer seed=${seed} → ${job}"
+done
diff --git a/scripts/slurm/submit_ehrmamba.sh b/scripts/slurm/submit_ehrmamba.sh
new file mode 100755
index 000000000..01fcca061
--- /dev/null
+++ b/scripts/slurm/submit_ehrmamba.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for ehrmamba — 3 fixed seeds, IllinoisComputes-GPU (A100/H200).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_ehrmamba.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_ehrmamba_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=12:00:00 \
+        --output="logs/slurm/table2_ehrmamba_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_ehrmamba_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=ehrmamba,SEED="${seed}",TABLE2_BS_EHRMAMBA=8 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted ehrmamba seed=${seed} → ${job}"
+done
diff --git a/scripts/slurm/submit_jambaehr.sh b/scripts/slurm/submit_jambaehr.sh
new file mode 100755
index 000000000..5a5429d4d
--- /dev/null
+++ b/scripts/slurm/submit_jambaehr.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for jambaehr — 3 fixed seeds, IllinoisComputes-GPU (A100/H200).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_jambaehr.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_jambaehr_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=18:00:00 \
+        --output="logs/slurm/table2_jambaehr_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_jambaehr_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=jambaehr,SEED="${seed}",TABLE2_BS_JAMBAEHR=4 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted jambaehr seed=${seed} → ${job}"
+done
diff --git a/scripts/slurm/submit_mlp.sh b/scripts/slurm/submit_mlp.sh
new file mode 100755
index 000000000..bd0265ec0
--- /dev/null
+++ b/scripts/slurm/submit_mlp.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for mlp — 3 fixed seeds, IllinoisComputes-GPU (A100/H200).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_mlp.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_mlp_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=6:00:00 \
+        --output="logs/slurm/table2_mlp_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_mlp_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=mlp,SEED="${seed}",TABLE2_BS_MLP=16 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted mlp seed=${seed} → ${job}"
+done
diff --git a/scripts/slurm/submit_rnn.sh b/scripts/slurm/submit_rnn.sh
new file mode 100755
index 000000000..900dbe51b
--- /dev/null
+++ b/scripts/slurm/submit_rnn.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for rnn — 3 fixed seeds, IllinoisComputes-GPU (A100/H200).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_rnn.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_rnn_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=6:00:00 \
+        --output="logs/slurm/table2_rnn_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_rnn_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=rnn,SEED="${seed}",TABLE2_BS_RNN=16 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted rnn seed=${seed} → ${job}"
+done
diff --git a/scripts/slurm/submit_table2_ic.sh b/scripts/slurm/submit_table2_ic.sh
new file mode 100755
index 000000000..1e23bccd5
--- /dev/null
+++ b/scripts/slurm/submit_table2_ic.sh
@@ -0,0 +1,100 @@
+#!/usr/bin/env bash
+# Submit 12 Table 2 jobs (heavy models) to IllinoisComputes-GPU partition.
+# Reads seeds from the existing manifest so results stay consistent with
+# mlp/rnn runs already queued on eng-research-gpu.
+#
+# Usage:
+#   bash scripts/slurm/submit_table2_ic.sh
+#
+# Before running: cancel the pending eng-research heavy model jobs:
+#   scancel <job_ids>   # ehrmamba/transformer/bottleneck/jambaehr jobs only
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+SEED_MANIFEST="${SEED_MANIFEST:-scripts/slurm/table2_random_seeds.txt}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+if [[ ! -f "${SEED_MANIFEST}" ]]; then
+    echo "ERROR: Seed manifest not found: ${SEED_MANIFEST}" >&2
+    echo "Run submit_table2_random.sh first to generate seeds, then rerun this script." >&2
+    exit 1
+fi
+
+mapfile -t SEEDS < "${SEED_MANIFEST}"
+if [[ "${#SEEDS[@]}" -ne 3 ]]; then
+    echo "ERROR: Expected 3 seeds in manifest, got ${#SEEDS[@]}" >&2
+    exit 1
+fi
+
+echo "Seeds (from manifest): ${SEEDS[*]}"
+echo "Partition : ${PARTITION}"
+echo "Account   : ${ACCOUNT}"
+echo ""
+
+submit_job() {
+    local model="$1"
+    local seed="$2"
+    local time_limit="$3"
+    local mem="$4"
+    shift 4
+    local extra_exports=("$@")   # additional KEY=VALUE strings
+
+    local export_str="ALL,MODEL=${model},SEED=${seed}"
+    for kv in "${extra_exports[@]}"; do
+        export_str="${export_str},${kv}"
+    done
+
+    sbatch \
+        --job-name="t2ic_${model}_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 \
+        --ntasks=1 \
+        --cpus-per-task=4 \
+        --mem="${mem}" \
+        --gres=gpu:1 \
+        --time="${time_limit}" \
+        --output="logs/slurm/table2ic_${model}_seed${seed}_%j.out" \
+        --error="logs/slurm/table2ic_${model}_seed${seed}_%j.err" \
+        --export="${export_str}" \
+        scripts/slurm/run_table2.sh
+}
+
+# A100 = 80GB VRAM, H200 = 141GB — can run much larger batches than A10 (24GB).
+# ehrmamba:              bs=8  (vs 2 on A10)
+# transformer:           bs=4  (vs 1 on A10)
+# bottleneck_transformer:bs=4  (vs 1 on A10)
+# jambaehr:              bs=4  (vs 1 on A10)
+
+echo "Submitting ehrmamba × 3..."
+for seed in "${SEEDS[@]}"; do
+    job=$(submit_job "ehrmamba" "${seed}" "12:00:00" "32G" "TABLE2_BS_EHRMAMBA=8")
+    echo "  ehrmamba seed=${seed} → ${job}"
+done
+
+echo "Submitting transformer × 3..."
+for seed in "${SEEDS[@]}"; do
+    job=$(submit_job "transformer" "${seed}" "18:00:00" "32G" "TABLE2_BS_TRANSFORMER=4")
+    echo "  transformer seed=${seed} → ${job}"
+done
+
+echo "Submitting bottleneck_transformer × 3..."
+for seed in "${SEEDS[@]}"; do
+    job=$(submit_job "bottleneck_transformer" "${seed}" "18:00:00" "32G" "TABLE2_BS_BOTTLENECK=4")
+    echo "  bottleneck_transformer seed=${seed} → ${job}"
+done
+
+echo "Submitting jambaehr × 3..."
+for seed in "${SEEDS[@]}"; do
+    job=$(submit_job "jambaehr" "${seed}" "18:00:00" "32G" "TABLE2_BS_JAMBAEHR=4")
+    echo "  jambaehr seed=${seed} → ${job}"
+done
+
+echo ""
+echo "12 IC jobs submitted. Monitor with:"
+echo "  squeue -u rianatri -p ${PARTITION}"
+echo "  tail -f logs/slurm/table2ic_<model>_seed<seed>_<jobid>.out"
diff --git a/scripts/slurm/submit_table2_random.sh b/scripts/slurm/submit_table2_random.sh
new file mode 100755
index 000000000..6a7902d06
--- /dev/null
+++ b/scripts/slurm/submit_table2_random.sh
@@ -0,0 +1,78 @@
+#!/usr/bin/env bash
+# Generates 3 random seeds, submits 18 SLURM jobs (1 per model+seed),
+# with heavier resource requests for memory-hungry models.
+# Optionally chains off a cachewarm job via --dependency=afterok:<jobid>.
+set -euo pipefail
+
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+SEED_MANIFEST="${SEED_MANIFEST:-scripts/slurm/table2_random_seeds.txt}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-cs-eng}"
+PARTITION="${SLURM_PARTITION:-eng-research-gpu}"
+CACHEWARM_JOB_ID="${1:-}"   # optional: pass cachewarm job ID to chain dependency
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+mapfile -t SEEDS < <(
+python3 - <<'PY'
+import random
+rng = random.SystemRandom()
+for seed in rng.sample(range(1, 2_147_483_647), 3):
+    print(seed)
+PY
+)
+
+if [[ "${#SEEDS[@]}" -ne 3 ]]; then
+    echo "ERROR: failed to generate 3 random seeds." >&2
+    exit 1
+fi
+
+printf '%s\n' "${SEEDS[@]}" > "${SEED_MANIFEST}"
+echo "Random seeds: ${SEEDS[*]}"
+echo "Manifest    : ${SEED_MANIFEST}"
+
+# Dependency string — only set if a cachewarm job ID was passed
+DEPEND=""
+if [[ -n "${CACHEWARM_JOB_ID}" ]]; then
+    DEPEND="--dependency=afterok:${CACHEWARM_JOB_ID}"
+    echo "Chaining after cachewarm job: ${CACHEWARM_JOB_ID}"
+fi
+
+submit_job() {
+    local model="$1"
+    local seed="$2"
+    local time_limit="$3"
+    local mem="$4"
+
+    sbatch \
+        --job-name="t2_${model}_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 \
+        --ntasks=1 \
+        --cpus-per-task=4 \
+        --mem="${mem}" \
+        --gres=gpu:1 \
+        --time="${time_limit}" \
+        --output="logs/slurm/table2_${model}_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_${model}_seed${seed}_%j.err" \
+        --export=ALL,MODEL="${model}",SEED="${seed}" \
+        ${DEPEND} \
+        scripts/slurm/run_table2.sh
+}
+
+# Light models (eng-research-gpu / A10 24GB) — mlp + rnn only.
+# ehrmamba and heavy models go to IllinoisComputes-GPU via submit_table2_ic.sh.
+for model in mlp rnn; do
+    for seed in "${SEEDS[@]}"; do
+        job=$(submit_job "${model}" "${seed}" "12:00:00" "32G")
+        echo "  Submitted ${model} seed=${seed} → ${job}"
+    done
+done
+
+echo ""
+echo "6 eng-research jobs submitted (mlp × 3, rnn × 3)."
+echo "Run submit_table2_ic.sh next for the remaining 12 (ehrmamba + heavy models)."
+echo "Monitor with:"
+echo "  squeue -u rianatri"
+echo "  tail -f logs/slurm/table2_<model>_seed<seed>_<jobid>.out"
diff --git a/scripts/slurm/submit_transformer.sh b/scripts/slurm/submit_transformer.sh
new file mode 100755
index 000000000..fef6329ff
--- /dev/null
+++ b/scripts/slurm/submit_transformer.sh
@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for transformer — 3 fixed seeds, IllinoisComputes-GPU (A100/H200).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/slurm/submit_transformer.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+ACCOUNT="${SLURM_ACCOUNT:-jimeng-ic}"
+PARTITION="${SLURM_PARTITION:-IllinoisComputes-GPU}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/slurm
+
+for seed in "${SEEDS[@]}"; do
+    job=$(sbatch \
+        --job-name="t2_transformer_s${seed}" \
+        --account="${ACCOUNT}" \
+        --partition="${PARTITION}" \
+        --nodes=1 --ntasks=1 --cpus-per-task=4 \
+        --mem=32G --gres=gpu:1 --time=18:00:00 \
+        --output="logs/slurm/table2_transformer_seed${seed}_%j.out" \
+        --error="logs/slurm/table2_transformer_seed${seed}_%j.err" \
+        --export=ALL,CACHE_DIR=/u/rianatri/pyhealth_cache,MODEL=transformer,SEED="${seed}",TABLE2_BS_TRANSFORMER=4 \
+        scripts/slurm/run_table2.sh | awk '{print $NF}')
+    echo "  Submitted transformer seed=${seed} → ${job}"
+done
diff --git a/scripts/sunlab/run_ehrmamba_timewindow.sh b/scripts/sunlab/run_ehrmamba_timewindow.sh
new file mode 100644
index 000000000..2ea666be9
--- /dev/null
+++ b/scripts/sunlab/run_ehrmamba_timewindow.sh
@@ -0,0 +1,36 @@
+#!/usr/bin/env bash
+# Run ehrmamba jobs for multiple seeds in tmux sessions
+# Usage: bash scripts/run_ehrmamba.sh
+set -euo pipefail
+
+SEEDS=(267573289 1872967241 706384748)
+TIME_WINDOWS=(24 48 96)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+USER=""
+
+cd "${PROJECT_DIR}"
+mkdir -p logs
+
+for seed in "${SEEDS[@]}"; do
+    for tw in "${TIME_WINDOWS[@]}"; do
+        echo "  Launching ehrmamba seed=${seed} time_window=${tw} in tmux"
+        tmux new -s "ehrmamba_${seed}_observationtimewindow_${tw}" "bash -lc '
+            conda activate pyhealth2 && \
+            CUDA_VISIBLE_DEVICES=4 python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \
+                --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \
+                --cache-dir /home/${USER}/pyhealth_cache \
+                --task icd_labs \
+                --model ehrmamba \
+                --embedding-dim 128 \
+                --num-layers 2 \
+                --observation-window-hours ${tw} \
+                --mamba-state-size 16 \
+                --mamba-conv-kernel 4 \
+                --epochs 20 \
+                --batch-size 8 \
+                --seed ${seed} \
+                --output-dir /home/${USER}/output/table2/ehrmamba_seed${seed}_observationtimewindow_${tw} \
+                2>&1 | tee logs/ehrmamba_seed${seed}_observationtimewindow_${tw}.log; exec bash'"
+        echo "  Launched ehrmamba seed=${seed} time_window=${tw} → tmux session: ehrmamba_${seed}_observationtimewindow_${tw}"
+    done
+done
diff --git a/scripts/sunlab/run_rnn.sh b/scripts/sunlab/run_rnn.sh
new file mode 100755
index 000000000..f47238805
--- /dev/null
+++ b/scripts/sunlab/run_rnn.sh
@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for rnn — 3 fixed seeds, sunlab cluster (tmux sessions).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/sunlab/run_rnn.sh
+set -euo pipefail
+
+# SEEDS=(267573289 1872967241 706384748)
+SEEDS=(42)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+USER="${USER:-wp14}"
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+CUDA_DEVICE="${CUDA_DEVICE:-0}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/home/${USER}/pyhealth_cache}"
+OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}"
+
+cd "${PROJECT_DIR}"
+mkdir -p logs/sunlab
+
+for seed in "${SEEDS[@]}"; do
+    session="rnn_seed${seed}"
+    echo "  Launching rnn seed=${seed} in tmux session: ${session}"
+    tmux new-session -d -s "${session}" "bash -lc '
+        conda activate ${CONDA_ENV} && \
+        CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \
+            --ehr-root ${EHR_ROOT} \
+            --cache-dir ${CACHE_DIR} \
+            --task icd_labs \
+            --model rnn \
+            --embedding-dim 128 \
+            --hidden-dim 128 \
+            --heads 4 \
+            --num-layers 2 \
+            --dropout 0.1 \
+            --epochs 20 \
+            --batch-size 16 \
+            --weight-decay 1e-5 \
+            --patience 5 \
+            --seed ${seed} \
+            --output-dir ${OUTPUT_BASE}/rnn_seed${seed} \
+            2>&1 | tee logs/sunlab/rnn_seed${seed}.log; exec bash'"
+    echo "  Launched rnn seed=${seed} → tmux session: ${session}"
+done
diff --git a/scripts/sunlab/run_transformer.sh b/scripts/sunlab/run_transformer.sh
new file mode 100644
index 000000000..0a2aacaff
--- /dev/null
+++ b/scripts/sunlab/run_transformer.sh
@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Submit Table 2 jobs for transformer — 3 fixed seeds, sunlab cluster (tmux sessions).
+# Seeds are shared across all lab members for aligned comparisons.
+# Usage: bash scripts/sunlab/run_transformer.sh
+set -euo pipefail
+
+# SEEDS=(267573289 1872967241 706384748)
+SEEDS=(44)
+PROJECT_DIR="${PROJECT_DIR:-$(pwd)}"
+USER="${USER:-wp14}"
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+CUDA_DEVICE="${CUDA_DEVICE:-4}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/home/${USER}/pyhealth_cache}"
+OUTPUT_BASE="${OUTPUT_BASE:-/home/${USER}/output/table2}"
+
+cd "${PROJECT_DIR}"
+
+for seed in "${SEEDS[@]}"; do
+    session="transformer_seed${seed}"
+    log_dir="logs/sunlab/$(date +%Y%m%d)"
+    mkdir -p "${log_dir}"
+    tmux new-session -d -s "${session}" "bash -lc '
+        conda activate ${CONDA_ENV} && \
+        CUDA_VISIBLE_DEVICES=${CUDA_DEVICE} python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \
+            --ehr-root ${EHR_ROOT} \
+            --cache-dir ${CACHE_DIR} \
+            --task icd_labs \
+            --model transformer \
+            --embedding-dim 64 \
+            --hidden-dim 64 \
+            --heads 2 \
+            --num-layers 1 \
+            --dropout 0.1 \
+            --epochs 20 \
+            --batch-size 4 \
+            --weight-decay 1e-5 \
+            --patience 5 \
+            --seed ${seed} \
+            --output-dir ${OUTPUT_BASE}/$(date +%Y%m%d) \
+            2>&1 | tee ${log_dir}/transformer_seed${seed}.log; exec bash'"
+    echo "  Launched transformer seed=${seed} → tmux session: ${session}, log: ${log_dir}/transformer_seed${seed}.log"
+done
diff --git a/scripts/train_unified.py b/scripts/train_unified.py
new file mode 100644
index 000000000..6d908b491
--- /dev/null
+++ b/scripts/train_unified.py
@@ -0,0 +1,163 @@
+"""Unified training CLI for PyHealth multimodal experiments.
+
+Reads a YAML config (with optional _inherit chain) then fires the E2E script
+with all resolved arguments.  CLI flags override YAML values.
+
+Usage:
+    python scripts/train_unified.py --config configs/train/e2e_baseline.yaml \
+        --model transformer --seed 0
+
+    # Smoke test
+    python scripts/train_unified.py --config configs/train/smoke.yaml
+
+    # Full run with frozen encoder
+    python scripts/train_unified.py --config configs/train/e2e_balanced.yaml \
+        --model ehrmamba --freeze-encoder --seed 42
+
+The script resolves _model_overrides from the YAML, merges CLI overrides, then
+calls unified_embedding_e2e_mimic4.py via subprocess so it can be used from
+condor jobs without reimplementing the full arg-build logic.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict
+
+
+def _load_yaml(path: Path) -> Dict[str, Any]:
+    try:
+        import yaml  # type: ignore
+    except ImportError:
+        raise SystemExit("PyYAML not installed. Run: pip install pyyaml")
+    with open(path) as f:
+        return yaml.safe_load(f) or {}
+
+
+def _resolve_config(config_path: Path) -> Dict[str, Any]:
+    """Load config and recursively merge _inherit chain."""
+    cfg = _load_yaml(config_path)
+    inherit = cfg.pop("_inherit", None)
+    if inherit:
+        parent_path = config_path.parent / inherit
+        parent = _resolve_config(parent_path)
+        # parent values are defaults; child values win
+        merged = {**parent, **cfg}
+        return merged
+    return cfg
+
+
+def _apply_model_overrides(cfg: Dict[str, Any]) -> Dict[str, Any]:
+    """Merge _model_overrides[model] into cfg if present."""
+    overrides = cfg.pop("_model_overrides", {})
+    model = cfg.get("model", "mlp")
+    if model in overrides:
+        cfg = {**cfg, **overrides[model]}
+    return cfg
+
+
+def _to_args(cfg: Dict[str, Any]) -> list[str]:
+    """Convert resolved config dict to CLI args for unified_embedding_e2e_mimic4.py."""
+    # Map config keys to CLI flags (underscores → hyphens, bool flags → store_true)
+    bool_flags = {
+        "icd_codes", "freeze_encoder", "include_vitals",
+        "balanced_sampling", "bidirectional",
+    }
+    skip_keys = {"_inherit", "_model_overrides"}
+
+    args = []
+    for key, val in cfg.items():
+        if key in skip_keys or val is None:
+            continue
+        flag = "--" + key.replace("_", "-")
+        if key in bool_flags:
+            if val:
+                args.append(flag)
+        else:
+            args += [flag, str(val)]
+    return args
+
+
+def _parse_cli() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Unified PyHealth training CLI — loads a YAML config then fires the E2E script."
+    )
+    parser.add_argument(
+        "--config",
+        required=True,
+        help="Path to a YAML config file (e.g. configs/train/e2e_baseline.yaml).",
+    )
+    # Pass-through overrides — any key from the YAML can be overridden here.
+    parser.add_argument("--model", default=None)
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--task", default=None)
+    parser.add_argument("--epochs", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--embedding-dim", type=int, default=None)
+    parser.add_argument("--dev", type=int, default=None)
+    parser.add_argument("--freeze-encoder", action="store_true", default=None)
+    parser.add_argument("--balanced-sampling", action="store_true", default=None)
+    parser.add_argument("--icd-codes", action="store_true", default=None)
+    parser.add_argument("--output-dir", default=None)
+    parser.add_argument("--ehr-root", default=None)
+    parser.add_argument("--note-root", default=None)
+    parser.add_argument("--cache-dir", default=None)
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Print the resolved command without running it.",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    cli = _parse_cli()
+
+    config_path = Path(cli.config)
+    if not config_path.exists():
+        raise SystemExit(f"Config not found: {config_path}")
+
+    cfg = _resolve_config(config_path)
+    cfg = _apply_model_overrides(cfg)
+
+    # Apply CLI overrides (non-None values win over YAML)
+    cli_overrides = {
+        "model": cli.model,
+        "seed": cli.seed,
+        "task": cli.task,
+        "epochs": cli.epochs,
+        "batch_size": cli.batch_size,
+        "embedding_dim": cli.embedding_dim,
+        "dev": cli.dev,
+        "freeze_encoder": cli.freeze_encoder if cli.freeze_encoder else None,
+        "balanced_sampling": cli.balanced_sampling if cli.balanced_sampling else None,
+        "icd_codes": cli.icd_codes if cli.icd_codes else None,
+        "output_dir": cli.output_dir,
+        "ehr_root": cli.ehr_root,
+        "note_root": cli.note_root,
+        "cache_dir": cli.cache_dir,
+    }
+    for k, v in cli_overrides.items():
+        if v is not None:
+            cfg[k] = v
+
+    script = Path(__file__).parent.parent / "examples" / "mortality_prediction" / "unified_embedding_e2e_mimic4.py"
+    cmd = [sys.executable, str(script)] + _to_args(cfg)
+
+    print("Resolved command:")
+    print("  " + " \\\n  ".join(cmd))
+
+    if cli.dry_run:
+        print("\n[dry-run] Not executing.")
+        return
+
+    result = subprocess.run(cmd, env={**os.environ, "PYTHONPATH": str(script.parent.parent.parent)})
+    sys.exit(result.returncode)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/will/compute_token_stats_labs_only.py b/scripts/will/compute_token_stats_labs_only.py
new file mode 100644
index 000000000..956f44444
--- /dev/null
+++ b/scripts/will/compute_token_stats_labs_only.py
@@ -0,0 +1,235 @@
+#!/usr/bin/env python3
+"""Compute per-modality token-count distributions and availability for Table 2.
+
+For each patient sample, computes a per-sample "token count" for every
+modality (notes: actual tokenizer token count; labs/ICD: number of real
+timesteps / codes), then reports the median and max of that distribution
+plus availability — the % of samples that have any real (non-placeholder)
+data for that modality.
+
+Iterates patients directly via dataset.iter_patients() and calls the task
+function on each one — never materializes all samples in memory, no litdata
+write, no GPU needed.
+
+Usage (on CC):
+    python scripts/compute_token_stats_labs_only.py \
+        --ehr-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2 \
+        --note-root /projects/illinois/eng/cs/jimeng/physionet.org/files/mimic-note \
+        --cache-dir /u/rianatri/pyhealth_cache \
+        --task notes_labs
+
+    python /home/wp14/PyHealth/scripts/compute_token_stats_labs_only.py \
+        --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \
+        --task labs
+
+    # Add --log-file to also save the printed stats to a file:
+    python PyHealth/scripts/compute_token_stats_labs_only.py \
+        --ehr-root /shared/rsaas/physionet.org/files/mimiciv/2.2 \
+        --task labs \
+        --log-file logs/token_stats_labs.log
+"""
+import argparse
+import os
+import sys
+
+import numpy as np
+from tqdm import tqdm
+
+
+class Tee:
+    """Duplicates writes to multiple streams (e.g. stdout + a log file)."""
+
+    def __init__(self, *streams):
+        self.streams = streams
+
+    def write(self, data):
+        for s in self.streams:
+            s.write(data)
+
+    def flush(self):
+        for s in self.streams:
+            s.flush()
+
+
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument("--ehr-root", required=True)
+    p.add_argument(
+        "--note-root",
+        default=None,
+        help="Required unless --task labs (labs-only task needs no notes).",
+    )
+    p.add_argument("--cache-dir", default=None)
+    p.add_argument("--dev", action="store_true", help="Limit to 1000 patients")
+    p.add_argument(
+        "--task",
+        type=str,
+        choices=["notes_labs", "labs"],
+        default="notes_labs",
+        help=(
+            "Task to profile. 'notes_labs' (default) uses admission-context "
+            "text sections; 'labs' profiles the labs-only baseline (no notes, "
+            "no note-root needed)."
+        ),
+    )
+    p.add_argument(
+        "--icd-codes",
+        action="store_true",
+        default=False,
+        help="Include discharge-coded ICD codes when using --task notes_labs.",
+    )
+    p.add_argument(
+        "--log-file",
+        default=None,
+        help="Path to also write stdout output to. If omitted, only prints to console.",
+    )
+    return p.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    if args.task != "labs" and not args.note_root:
+        raise SystemExit("--note-root is required unless --task labs")
+
+    if args.log_file:
+        log_f = open(args.log_file, "w")
+        sys.stdout = Tee(sys.stdout, log_f)
+        sys.stderr = Tee(sys.stderr, log_f)
+        print(f"Logging output to {args.log_file}")
+
+    os.environ.setdefault("PYHEALTH_DISABLE_DASK_DISTRIBUTED", "1")
+
+    from pyhealth.datasets import MIMIC4Dataset
+    from pyhealth.tasks.multimodal_mimic4 import LabsMIMIC4, NotesLabsMIMIC4
+
+    print("Building dataset (uses cache if available)...")
+
+    if args.task == "labs":
+        ehr_tables = ["labevents"]
+        note_tables = []
+        task = LabsMIMIC4(window_hours=24)
+    else:  # notes_labs
+        ehr_tables = (
+            ["diagnoses_icd", "procedures_icd", "labevents"]
+            if args.icd_codes
+            else ["labevents"]
+        )
+        note_tables = ["discharge"]
+        task = NotesLabsMIMIC4(window_hours=24, include_icd=args.icd_codes)
+
+    kwargs = dict(
+        ehr_root=args.ehr_root,
+        ehr_tables=ehr_tables,
+        note_tables=note_tables,
+        dev=args.dev,
+    )
+    if args.note_root:
+        kwargs["note_root"] = args.note_root
+    if args.cache_dir:
+        kwargs["cache_dir"] = args.cache_dir
+
+    dataset = MIMIC4Dataset(**kwargs)
+
+    MISSING_TEXT = ""
+    has_notes = args.task == "notes_labs"
+    has_icd = args.task == "notes_labs" and args.icd_codes
+
+    tokenizer = None
+    if has_notes:
+        from transformers import AutoTokenizer
+
+        tokenizer_model = task.input_schema["admission_note_times"][1][
+            "tokenizer_model"
+        ]
+        print(f"Loading tokenizer: {tokenizer_model}")
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_model)
+
+    # Per-sample token counts (one entry per patient sample)
+    note_token_counts = []
+    note_has_data = []
+    lab_token_counts = []
+    lab_has_data = []
+    icd_token_counts = []
+    icd_has_data = []
+
+    n_patients = n_samples = 0
+
+    print("Iterating patients and accumulating stats (no litdata write)...")
+    for patient in tqdm(dataset.iter_patients(), total=len(dataset.unique_patient_ids)):
+        n_patients += 1
+        samples = task(patient)
+        if not samples:
+            continue
+        for s in samples:
+            n_samples += 1
+
+            # ── notes ────────────────────────────────────────────
+            if has_notes:
+                note_texts, _ = s["admission_note_times"]
+                real_texts = [t for t in note_texts if t != MISSING_TEXT]
+                n_tokens = sum(
+                    len(tokenizer.encode(t, add_special_tokens=False))
+                    for t in real_texts
+                )
+                note_token_counts.append(n_tokens)
+                note_has_data.append(bool(real_texts))
+
+            # ── ICD codes ────────────────────────────────────────
+            if has_icd:
+                _, icd_visits = s["icd_codes"]
+                real_visits = [v for v in icd_visits if v != [MISSING_TEXT]]
+                n_codes = sum(len(v) for v in real_visits)
+                icd_token_counts.append(n_codes)
+                icd_has_data.append(bool(real_visits))
+
+            # ── labs ─────────────────────────────────────────────
+            _, lab_masks = s["labs_mask"]
+            n_real_timesteps = sum(1 for mask_row in lab_masks if any(mask_row))
+            lab_token_counts.append(n_real_timesteps)
+            lab_has_data.append(n_real_timesteps > 0)
+
+    def median_max(counts):
+        if not counts:
+            return float("nan"), float("nan")
+        arr = np.array(counts)
+        return float(np.median(arr)), float(np.max(arr))
+
+    def availability_pct(has_data):
+        if not has_data:
+            return float("nan")
+        return 100.0 * sum(has_data) / len(has_data)
+
+    print("\n" + "=" * 60)
+    print(f"TOKEN STATS — {task.task_name}")
+    print("=" * 60)
+    print(f"  Patients processed : {n_patients:>8,}")
+    print(f"  Samples (patients) : {n_samples:>8,}")
+
+    if has_notes:
+        med, mx = median_max(note_token_counts)
+        avail = availability_pct(note_has_data)
+        print("\n── Notes (admission-context) ───────────────────────────────")
+        print(f"  Median tokens : {med:>10.1f}")
+        print(f"  Max tokens    : {mx:>10.0f}")
+        print(f"  Availability  : {avail:>9.1f}%")
+
+    if has_icd:
+        med, mx = median_max(icd_token_counts)
+        avail = availability_pct(icd_has_data)
+        print("\n── ICD Codes ────────────────────────────────────────────────")
+        print(f"  Median tokens : {med:>10.1f}")
+        print(f"  Max tokens    : {mx:>10.0f}")
+        print(f"  Availability  : {avail:>9.1f}%")
+
+    med, mx = median_max(lab_token_counts)
+    avail = availability_pct(lab_has_data)
+    print("\n── Labs ─────────────────────────────────────────────────────")
+    print(f"  Median tokens (timesteps) : {med:>10.1f}")
+    print(f"  Max tokens                : {mx:>10.0f}")
+    print(f"  Availability              : {avail:>9.1f}%")
+    print()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/will/condor/labs_notes/labs_notes_rnn.sub b/scripts/will/condor/labs_notes/labs_notes_rnn.sub
new file mode 100644
index 000000000..f8e7a1ef2
--- /dev/null
+++ b/scripts/will/condor/labs_notes/labs_notes_rnn.sub
@@ -0,0 +1,59 @@
+# HTCondor submission — labs+notes RNN mortality run
+#
+# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py:
+# same task (notes_labs), same model (rnn), same hyperparameters. Runs unattended
+# instead of in a tmux session; Condor assigns the GPU (no manual
+# nvidia-smi / CUDA_VISIBLE_DEVICES step needed).
+#
+# To submit (from the project root):
+#   mkdir -p /home/wp14/logs/condor
+#   condor_submit scripts/will/condor/labs_notes/labs_notes_rnn.sub
+#
+# Monitor:
+#   condor_q
+#   tail -f /home/wp14/logs/condor/labs_notes_rnn_<ClusterId>_0.out
+
+initialdir          = /home/wp14/PyHealth
+executable          = /home/wp14/PyHealth/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh
+transfer_executable  = False
+arguments            = $(seed)
+getenv               = True
+
+environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \
+               NOTE_ROOT=/shared/rsaas/physionet.org/files/mimic-note \
+               CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes \
+               OUTPUT_DIR=/home/wp14/output \
+               USE_WANDB=1 \
+               WANDB_PROJECT=pyhealth-multimodal-labs-notes \
+               WANDB_RUN_NAME=labs_notes_rnn_seed$(seed) \
+               FREEZE_ENCODER=1"
+
+output = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).out
+error  = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).err
+log    = /home/wp14/logs/condor/labs_notes_rnn_$(ClusterId)_$(Process).log
+
+stream_output = False
+stream_error  = False
+
+request_gpus   = 1
+request_cpus   = 4
+# Was 163840MB (160GB) — cgroup-killed job 10821 at 162747MB during the
+# full-scale patient_id sort/shuffle (see run_labs_notes_rnn.sh comment).
+# Bumped for headroom now that the distributed cluster (with disk-spilling)
+# is back in play; c02 has ~1TB total and is otherwise idle.
+request_memory = 400000MB
+request_disk   = 20GB
+
+# Previously hardcoded to sunlab-c01 (A100 80GB) because the previous run
+# OOM'd on a 47GB card with ~46.8GB resident. sunlab-c01's condor_startd is
+# currently down (master alive, STARTD_StartTime=0), so it never matches —
+# the only machine in the pool is sunlab-c02 (8x RTX 6000 Ada, 48509MB each).
+# Match any GPU with enough headroom and prefer the biggest; FREEZE_ENCODER=1
+# below (frozen Bio_ClinicalBERT text encoder, ~50% less VRAM for the text
+# branch) is what actually keeps this under 48GB instead of the hostname pin.
+Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000)
+Rank = TARGET.GPUs_GlobalMemoryMb
+
+queue seed from (
+    12
+)
diff --git a/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh
new file mode 100755
index 000000000..c75feda7f
--- /dev/null
+++ b/scripts/will/condor/labs_notes/run_labs_notes_rnn.sh
@@ -0,0 +1,193 @@
+#!/usr/bin/env bash
+# HTCondor executable — labs+notes RNN mortality run.
+#
+# Condor equivalent of scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py:
+# same task (notes_labs), same model (rnn), same hyperparameter defaults. GPU
+# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor
+# assigns the GPU via request_gpus / cgroups.
+#
+# usage: run_labs_notes_rnn.sh <seed>
+# to remove logs: rm -rf logs/condor/*
+# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes/*
+set -euo pipefail
+
+SEED="${1:?usage: run_labs_notes_rnn.sh <seed>}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}"
+CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs_notes}"
+OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}"
+CONDA_SH="${CONDA_SH:-}"
+
+DEV_MODE="${DEV_MODE:-0}"
+EMBEDDING_DIM="${EMBEDDING_DIM:-128}"
+HIDDEN_DIM="${HIDDEN_DIM:-128}"
+RNN_TYPE="${RNN_TYPE:-GRU}"
+RNN_LAYERS="${RNN_LAYERS:-2}"
+DROPOUT="${DROPOUT:-0.1}"
+EPOCHS="${EPOCHS:-15}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+LR="${LR:-1e-3}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}"
+PATIENCE="${PATIENCE:-5}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+FREEZE_ENCODER="${FREEZE_ENCODER:-0}"
+INCLUDE_VITALS="${INCLUDE_VITALS:-0}"
+POS_WEIGHT="${POS_WEIGHT:-1}"
+USE_AMP="${USE_AMP:-0}"
+AMP_DTYPE="${AMP_DTYPE:-bf16}"
+
+# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that
+# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup.
+# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1)
+# also throws away its disk-spilling/memory limits, which OOM'd the full-scale
+# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup
+# limit hit). Scope the fix to just the NVML probe instead, keeping the real
+# distributed cluster (with spilling) for the sort.
+export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}"
+
+USE_WANDB="${USE_WANDB:-0}"
+WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes}"
+WANDB_RUN_NAME="${WANDB_RUN_NAME:-}"
+
+resolve_conda_sh() {
+    if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [[ -f "${c}" ]]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [[ -n "${found}" && -f "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then
+    echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2
+    exit 1
+fi
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+
+JOB_TAG="rnn_labs_notes_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}"
+
+echo "========================================================"
+echo "  Labs+notes RNN run  |  ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Conda env : ${CONDA_ENV}"
+echo "  EHR root  : ${EHR_ROOT}"
+echo "  Note root : ${NOTE_ROOT}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  Output dir: ${OUTPUT_DIR}"
+echo "  Seed      : ${SEED}"
+echo "  Dev mode  : ${DEV_MODE}"
+echo "  Use AMP   : ${USE_AMP} (dtype=${AMP_DTYPE})"
+echo "========================================================"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2
+    exit 1
+fi
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --note-root "${NOTE_ROOT}"
+    --cache-dir "${CACHE_DIR}"
+    --task notes_labs
+    --model rnn
+    --embedding-dim "${EMBEDDING_DIM}"
+    --hidden-dim "${HIDDEN_DIM}"
+    --rnn-type "${RNN_TYPE}"
+    --rnn-layers "${RNN_LAYERS}"
+    --dropout "${DROPOUT}"
+    --epochs "${EPOCHS}"
+    --batch-size "${BATCH_SIZE}"
+    --lr "${LR}"
+    --weight-decay "${WEIGHT_DECAY}"
+    --patience "${PATIENCE}"
+    --num-workers "${NUM_WORKERS}"
+    --seed "${SEED}"
+    --output-dir "${OUTPUT_DIR}"
+    --pos-weight "${POS_WEIGHT}"
+)
+
+if [[ "${DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev)
+fi
+
+if [[ "${FREEZE_ENCODER}" == "1" ]]; then
+    COMMON+=(--freeze-encoder)
+fi
+
+if [[ "${INCLUDE_VITALS}" == "1" ]]; then
+    COMMON+=(--include-vitals)
+fi
+
+if [[ "${USE_AMP}" == "1" ]]; then
+    COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}")
+fi
+
+if [[ "${USE_WANDB}" == "1" ]]; then
+    COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}")
+    if [[ -n "${WANDB_RUN_NAME}" ]]; then
+        COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}")
+    fi
+fi
+
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}"
+
+echo "========================================================"
+echo "  Completed ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub
new file mode 100644
index 000000000..118fa8149
--- /dev/null
+++ b/scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub
@@ -0,0 +1,61 @@
+# HTCondor submission — labs+notes+CXR RNN mortality run
+#
+# Condor equivalent of
+# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py:
+# same task (notes_labs_cxr), same model (rnn), same hyperparameters. Runs
+# unattended instead of in a tmux session; Condor assigns the GPU (no manual
+# nvidia-smi / CUDA_VISIBLE_DEVICES step needed).
+#
+# To submit (from the project root):
+#   mkdir -p /home/wp14/logs/condor
+#   condor_submit scripts/will/condor/labs_notes_cxr/labs_notes_cxr_rnn.sub
+#
+# Monitor:
+#   condor_q
+#   tail -f /home/wp14/logs/condor/labs_notes_cxr_rnn_<ClusterId>_0.out
+
+initialdir          = /home/wp14/PyHealth
+executable          = /home/wp14/PyHealth/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh
+transfer_executable  = False
+arguments            = $(seed)
+getenv               = True
+
+environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \
+               NOTE_ROOT=/shared/rsaas/physionet.org/files/mimic-note \
+               CXR_ROOT=/shared/rsaas/physionet.org/files/MIMIC-CXR \
+               CXR_VARIANT=sunlab \
+               CACHE_DIR=/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr \
+               OUTPUT_DIR=/home/wp14/output \
+               USE_WANDB=1 \
+               WANDB_PROJECT=pyhealth-multimodal-labs-notes-cxr \
+               WANDB_RUN_NAME=labs_notes_cxr_rnn_seed$(seed) \
+               FREEZE_ENCODER=1"
+
+output = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).out
+error  = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).err
+log    = /home/wp14/logs/condor/labs_notes_cxr_rnn_$(ClusterId)_$(Process).log
+
+stream_output = False
+stream_error  = False
+
+request_gpus   = 1
+request_cpus   = 8
+# Starting point copied from labs_notes_rnn.sub (400000MB was sized for the
+# full-scale patient_id sort/shuffle there). Untested for this variant: the
+# added CXR branch means more dataloader workers decoding JPEGs plus an
+# image-metadata join during caching, so this may need to be raised further
+# if the job gets cgroup-killed.
+request_memory = 400000MB
+request_disk   = 20GB
+
+# Same reasoning as labs_notes_rnn.sub: FREEZE_ENCODER=1 (frozen
+# Bio_ClinicalBERT text encoder) keeps VRAM down for the text branch, but the
+# CXR image encoder is unfrozen here and adds its own footprint on top, so
+# the GPU memory floor is a conservative starting guess pending a real OOM
+# data point on this variant.
+Requirements = (TARGET.GPUs_GlobalMemoryMb >= 40000)
+Rank = TARGET.GPUs_GlobalMemoryMb
+
+queue seed from (
+    12
+)
diff --git a/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh
new file mode 100755
index 000000000..82a16b85a
--- /dev/null
+++ b/scripts/will/condor/labs_notes_cxr/run_labs_notes_cxr_rnn.sh
@@ -0,0 +1,199 @@
+#!/usr/bin/env bash
+# HTCondor executable — labs+notes+CXR RNN mortality run.
+#
+# Condor equivalent of
+# scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py:
+# same task (notes_labs_cxr), same model (rnn), same hyperparameter defaults. GPU
+# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor
+# assigns the GPU via request_gpus / cgroups.
+#
+# usage: run_labs_notes_cxr_rnn.sh <seed>
+# to remove logs: rm -rf logs/condor/*
+# to remove cache: rm -rf /shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr/*
+set -euo pipefail
+
+SEED="${1:?usage: run_labs_notes_cxr_rnn.sh <seed>}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+NOTE_ROOT="${NOTE_ROOT:-/shared/rsaas/physionet.org/files/mimic-note}"
+CXR_ROOT="${CXR_ROOT:-/shared/rsaas/physionet.org/files/MIMIC-CXR}"
+CXR_VARIANT="${CXR_VARIANT:-default}"
+CACHE_DIR="${CACHE_DIR:-/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr}"
+OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}"
+CONDA_SH="${CONDA_SH:-}"
+
+DEV_MODE="${DEV_MODE:-0}"
+EMBEDDING_DIM="${EMBEDDING_DIM:-128}"
+HIDDEN_DIM="${HIDDEN_DIM:-128}"
+RNN_TYPE="${RNN_TYPE:-GRU}"
+RNN_LAYERS="${RNN_LAYERS:-2}"
+DROPOUT="${DROPOUT:-0.1}"
+EPOCHS="${EPOCHS:-15}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+LR="${LR:-1e-3}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}"
+PATIENCE="${PATIENCE:-5}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+FREEZE_ENCODER="${FREEZE_ENCODER:-0}"
+INCLUDE_VITALS="${INCLUDE_VITALS:-0}"
+POS_WEIGHT="${POS_WEIGHT:-1}"
+USE_AMP="${USE_AMP:-0}"
+AMP_DTYPE="${AMP_DTYPE:-bf16}"
+
+# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that
+# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup.
+# Disabling the whole distributed cluster (PYHEALTH_DISABLE_DASK_DISTRIBUTED=1)
+# also throws away its disk-spilling/memory limits, which OOM'd the full-scale
+# patient_id sort during event-dataframe caching (job 10821, 160GB cgroup
+# limit hit). Scope the fix to just the NVML probe instead, keeping the real
+# distributed cluster (with spilling) for the sort.
+export DASK_DISTRIBUTED__DIAGNOSTICS__NVML="${DASK_DISTRIBUTED__DIAGNOSTICS__NVML:-0}"
+
+USE_WANDB="${USE_WANDB:-0}"
+WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-notes-cxr}"
+WANDB_RUN_NAME="${WANDB_RUN_NAME:-}"
+
+resolve_conda_sh() {
+    if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [[ -f "${c}" ]]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [[ -n "${found}" && -f "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then
+    echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2
+    exit 1
+fi
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+
+JOB_TAG="rnn_labs_notes_cxr_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}"
+
+echo "========================================================"
+echo "  Labs+notes+CXR RNN run  |  ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Conda env  : ${CONDA_ENV}"
+echo "  EHR root   : ${EHR_ROOT}"
+echo "  Note root  : ${NOTE_ROOT}"
+echo "  CXR root   : ${CXR_ROOT} (variant=${CXR_VARIANT})"
+echo "  Cache dir  : ${CACHE_DIR}"
+echo "  Output dir : ${OUTPUT_DIR}"
+echo "  Seed       : ${SEED}"
+echo "  Dev mode   : ${DEV_MODE}"
+echo "  Use AMP    : ${USE_AMP} (dtype=${AMP_DTYPE})"
+echo "========================================================"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2
+    exit 1
+fi
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --note-root "${NOTE_ROOT}"
+    --cxr-root "${CXR_ROOT}"
+    --cxr-variant "${CXR_VARIANT}"
+    --cache-dir "${CACHE_DIR}"
+    --task notes_labs_cxr
+    --model rnn
+    --embedding-dim "${EMBEDDING_DIM}"
+    --hidden-dim "${HIDDEN_DIM}"
+    --rnn-type "${RNN_TYPE}"
+    --rnn-layers "${RNN_LAYERS}"
+    --dropout "${DROPOUT}"
+    --epochs "${EPOCHS}"
+    --batch-size "${BATCH_SIZE}"
+    --lr "${LR}"
+    --weight-decay "${WEIGHT_DECAY}"
+    --patience "${PATIENCE}"
+    --num-workers "${NUM_WORKERS}"
+    --seed "${SEED}"
+    --output-dir "${OUTPUT_DIR}"
+    --pos-weight "${POS_WEIGHT}"
+)
+
+if [[ "${DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev)
+fi
+
+if [[ "${FREEZE_ENCODER}" == "1" ]]; then
+    COMMON+=(--freeze-encoder)
+fi
+
+if [[ "${INCLUDE_VITALS}" == "1" ]]; then
+    COMMON+=(--include-vitals)
+fi
+
+if [[ "${USE_AMP}" == "1" ]]; then
+    COMMON+=(--use-amp --amp-dtype "${AMP_DTYPE}")
+fi
+
+if [[ "${USE_WANDB}" == "1" ]]; then
+    COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}")
+    if [[ -n "${WANDB_RUN_NAME}" ]]; then
+        COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}")
+    fi
+fi
+
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}"
+
+echo "========================================================"
+echo "  Completed ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/will/condor/labs_only/labs_only_rnn.sub b/scripts/will/condor/labs_only/labs_only_rnn.sub
new file mode 100644
index 000000000..045743401
--- /dev/null
+++ b/scripts/will/condor/labs_only/labs_only_rnn.sub
@@ -0,0 +1,45 @@
+# HTCondor submission — labs-only RNN mortality run
+#
+# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py:
+# same task (labs), same model (rnn), same hyperparameters. Runs unattended
+# instead of in a tmux session; Condor assigns the GPU (no manual
+# nvidia-smi / CUDA_VISIBLE_DEVICES step needed).
+#
+# To submit (from the project root):
+#   mkdir -p /home/wp14/logs/condor
+#   condor_submit scripts/will/condor/labs_only_rnn.sub
+#
+# Monitor:
+#   condor_q
+#   tail -f /home/wp14/logs/condor/labs_only_rnn_<ClusterId>_0.out
+
+initialdir          = /home/wp14/PyHealth
+executable          = /home/wp14/PyHealth/scripts/will/condor/run_labs_only_rnn.sh
+transfer_executable  = False
+arguments            = $(seed)
+getenv               = True
+
+environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \
+               CACHE_DIR=/shared/eng/wp14/pyhealth_cache_labs \
+               OUTPUT_DIR=/home/wp14/output \
+               USE_WANDB=1 \
+               WANDB_PROJECT=pyhealth-multimodal-labs-only \
+               WANDB_RUN_NAME=labs_rnn_seed$(seed)"
+
+output = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).out
+error  = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).err
+log    = /home/wp14/logs/condor/labs_only_rnn_$(ClusterId)_$(Process).log
+
+stream_output = False
+stream_error  = False
+
+request_gpus   = 1
+request_cpus   = 4
+request_memory = 163840MB
+request_disk   = 20GB
+
+Rank = TARGET.CUDAGlobalMemoryMb
+
+queue seed from (
+    12
+)
diff --git a/scripts/will/condor/labs_only/labs_only_transformer.sub b/scripts/will/condor/labs_only/labs_only_transformer.sub
new file mode 100644
index 000000000..99c1737c9
--- /dev/null
+++ b/scripts/will/condor/labs_only/labs_only_transformer.sub
@@ -0,0 +1,45 @@
+# HTCondor submission — labs-only Transformer mortality run
+#
+# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py:
+# same task (labs), transformer model, same hyperparameter pattern as
+# labs_only_rnn.sub. Runs unattended instead of in a tmux session; Condor
+# assigns the GPU (no manual nvidia-smi / CUDA_VISIBLE_DEVICES step needed).
+#
+# To submit (from the project root):
+#   mkdir -p /home/wp14/logs/condor
+#   condor_submit scripts/will/condor/labs_only/labs_only_transformer.sub
+#
+# Monitor:
+#   condor_q
+#   tail -f /home/wp14/logs/condor/labs_only_transformer_<ClusterId>_0.out
+
+initialdir          = /home/wp14/PyHealth
+executable          = /home/wp14/PyHealth/scripts/will/condor/labs_only/run_labs_only_transformer.sh
+transfer_executable  = False
+arguments            = $(seed)
+getenv               = True
+
+environment = "EHR_ROOT=/shared/rsaas/physionet.org/files/mimiciv/2.2 \
+               CACHE_DIR=/shared/eng/wp14/pyhealth_cache_labs \
+               OUTPUT_DIR=/home/wp14/output \
+               USE_WANDB=1 \
+               WANDB_PROJECT=pyhealth-multimodal-labs-only \
+               WANDB_RUN_NAME=labs_transformer_seed$(seed)"
+
+output = /home/wp14/logs/condor/labs_only_transformer_$(ClusterId)_$(Process).out
+error  = /home/wp14/logs/condor/labs_only_transformer_$(ClusterId)_$(Process).err
+log    = /home/wp14/logs/condor/labs_only_transformer_$(ClusterId)_$(Process).log
+
+stream_output = False
+stream_error  = False
+
+request_gpus   = 1
+request_cpus   = 4
+request_memory = 163840MB
+request_disk   = 20GB
+
+Rank = TARGET.CUDAGlobalMemoryMb
+
+queue seed from (
+    12
+)
diff --git a/scripts/will/condor/labs_only/run_labs_only_rnn.sh b/scripts/will/condor/labs_only/run_labs_only_rnn.sh
new file mode 100755
index 000000000..a9075b152
--- /dev/null
+++ b/scripts/will/condor/labs_only/run_labs_only_rnn.sh
@@ -0,0 +1,170 @@
+#!/usr/bin/env bash
+# HTCondor executable — labs-only RNN mortality run.
+#
+# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py:
+# same task (labs), same model (rnn), same hyperparameter defaults. GPU
+# selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is dropped since Condor
+# assigns the GPU via request_gpus / cgroups.
+#
+# usage: run_labs_only_rnn.sh <seed>
+# to remove logs: rm -rf logs/condor/*
+# to remove cache: rm -rf /shared/eng/wp14/pyhealth_cache_labs/*
+set -euo pipefail
+
+SEED="${1:?usage: run_labs_only_rnn.sh <seed>}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs}"
+OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}"
+CONDA_SH="${CONDA_SH:-}"
+
+DEV_MODE="${DEV_MODE:-0}"
+EMBEDDING_DIM="${EMBEDDING_DIM:-64}"
+HIDDEN_DIM="${HIDDEN_DIM:-64}"
+RNN_TYPE="${RNN_TYPE:-GRU}"
+RNN_LAYERS="${RNN_LAYERS:-1}"
+DROPOUT="${DROPOUT:-0.1}"
+EPOCHS="${EPOCHS:-15}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+LR="${LR:-1e-3}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}"
+PATIENCE="${PATIENCE:-5}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+POS_WEIGHT="${POS_WEIGHT:-1.0}"
+
+# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that
+# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup.
+# Skip the distributed Dask cluster (falls back to the plain local scheduler)
+# to avoid touching NVML during event-dataframe preprocessing.
+export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}"
+
+USE_WANDB="${USE_WANDB:-0}"
+WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-only}"
+WANDB_RUN_NAME="${WANDB_RUN_NAME:-}"
+
+resolve_conda_sh() {
+    if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [[ -f "${c}" ]]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [[ -n "${found}" && -f "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then
+    echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2
+    exit 1
+fi
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+
+JOB_TAG="rnn_labs_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}"
+
+echo "========================================================"
+echo "  Labs-only RNN run  |  ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Conda env : ${CONDA_ENV}"
+echo "  EHR root  : ${EHR_ROOT}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  Output dir: ${OUTPUT_DIR}"
+echo "  Seed      : ${SEED}"
+echo "  Dev mode  : ${DEV_MODE}"
+echo "========================================================"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2
+    exit 1
+fi
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --cache-dir "${CACHE_DIR}"
+    --task labs
+    --model rnn
+    --embedding-dim "${EMBEDDING_DIM}"
+    --hidden-dim "${HIDDEN_DIM}"
+    --rnn-type "${RNN_TYPE}"
+    --rnn-layers "${RNN_LAYERS}"
+    --dropout "${DROPOUT}"
+    --epochs "${EPOCHS}"
+    --batch-size "${BATCH_SIZE}"
+    --lr "${LR}"
+    --weight-decay "${WEIGHT_DECAY}"
+    --patience "${PATIENCE}"
+    --num-workers "${NUM_WORKERS}"
+    --pos-weight "${POS_WEIGHT}"
+    --seed "${SEED}"
+    --output-dir "${OUTPUT_DIR}"
+)
+
+if [[ "${DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev)
+fi
+
+if [[ "${USE_WANDB}" == "1" ]]; then
+    COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}")
+    if [[ -n "${WANDB_RUN_NAME}" ]]; then
+        COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}")
+    fi
+fi
+
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}"
+
+echo "========================================================"
+echo "  Completed ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/will/condor/labs_only/run_labs_only_transformer.sh b/scripts/will/condor/labs_only/run_labs_only_transformer.sh
new file mode 100755
index 000000000..972b5dc24
--- /dev/null
+++ b/scripts/will/condor/labs_only/run_labs_only_transformer.sh
@@ -0,0 +1,168 @@
+#!/usr/bin/env bash
+# HTCondor executable — labs-only Transformer mortality run.
+#
+# Condor equivalent of scripts/will/sunlab/tmux_run_labs_only_rnn_variant.py:
+# same task (labs), same defaults pattern as run_labs_only_rnn.sh but for the
+# transformer model. GPU selection (nvidia-smi / CUDA_VISIBLE_DEVICES) is
+# dropped since Condor assigns the GPU via request_gpus / cgroups.
+#
+# usage: run_labs_only_transformer.sh <seed>
+# to remove logs: rm -rf logs/condor/*
+# to remove cache: rm -rf /shared/eng/wp14/pyhealth_cache_labs/*
+set -euo pipefail
+
+SEED="${1:?usage: run_labs_only_transformer.sh <seed>}"
+
+CONDA_ENV="${CONDA_ENV:-pyhealth2}"
+PROJECT_DIR="${PROJECT_DIR:-/home/wp14/PyHealth}"
+EHR_ROOT="${EHR_ROOT:-/shared/rsaas/physionet.org/files/mimiciv/2.2}"
+CACHE_DIR="${CACHE_DIR:-/shared/eng/wp14/pyhealth_cache_labs}"
+OUTPUT_DIR="${OUTPUT_DIR:-/home/wp14/output}"
+CONDA_SH="${CONDA_SH:-}"
+
+DEV_MODE="${DEV_MODE:-0}"
+EMBEDDING_DIM="${EMBEDDING_DIM:-64}"
+HEADS="${HEADS:-4}"
+NUM_LAYERS="${NUM_LAYERS:-2}"
+DROPOUT="${DROPOUT:-0.1}"
+EPOCHS="${EPOCHS:-15}"
+BATCH_SIZE="${BATCH_SIZE:-32}"
+LR="${LR:-1e-3}"
+WEIGHT_DECAY="${WEIGHT_DECAY:-1e-5}"
+PATIENCE="${PATIENCE:-5}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+POS_WEIGHT="${POS_WEIGHT:-1.0}"
+
+# Condor GPU cgroups can expose a truncated CUDA_VISIBLE_DEVICES UUID that
+# distributed's NVML diagnostics can't resolve, crashing LocalCluster startup.
+# Skip the distributed Dask cluster (falls back to the plain local scheduler)
+# to avoid touching NVML during event-dataframe preprocessing.
+export PYHEALTH_DISABLE_DASK_DISTRIBUTED="${PYHEALTH_DISABLE_DASK_DISTRIBUTED:-1}"
+
+USE_WANDB="${USE_WANDB:-0}"
+WANDB_PROJECT="${WANDB_PROJECT:-pyhealth-multimodal-labs-only}"
+WANDB_RUN_NAME="${WANDB_RUN_NAME:-}"
+
+resolve_conda_sh() {
+    if [[ -n "${CONDA_SH}" && -f "${CONDA_SH}" ]]; then
+        echo "${CONDA_SH}"
+        return 0
+    fi
+    if command -v conda >/dev/null 2>&1; then
+        local base
+        base="$(conda info --base 2>/dev/null || true)"
+        if [[ -n "${base}" && -f "${base}/etc/profile.d/conda.sh" ]]; then
+            echo "${base}/etc/profile.d/conda.sh"
+            return 0
+        fi
+    fi
+    if [ -f /etc/profile.d/modules.sh ]; then
+        source /etc/profile.d/modules.sh >/dev/null 2>&1 || true
+        if command -v module >/dev/null 2>&1; then
+            module load miniconda3 >/dev/null 2>&1 || true
+            module load anaconda3 >/dev/null 2>&1 || true
+            if command -v conda >/dev/null 2>&1; then
+                local mod_base
+                mod_base="$(conda info --base 2>/dev/null || true)"
+                if [[ -n "${mod_base}" && -f "${mod_base}/etc/profile.d/conda.sh" ]]; then
+                    echo "${mod_base}/etc/profile.d/conda.sh"
+                    return 0
+                fi
+            fi
+        fi
+    fi
+    local user_name home_dir
+    user_name="${USER:-$(id -un 2>/dev/null || true)}"
+    home_dir="${HOME:-/home/${user_name}}"
+    local candidates=(
+        "${home_dir}/miniconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/miniconda3/etc/profile.d/conda.sh"
+        "${home_dir}/anaconda3/etc/profile.d/conda.sh"
+        "/home/${user_name}/anaconda3/etc/profile.d/conda.sh"
+        "/opt/miniconda3/etc/profile.d/conda.sh"
+        "/opt/anaconda3/etc/profile.d/conda.sh"
+        "/opt/conda/etc/profile.d/conda.sh"
+    )
+    local c
+    for c in "${candidates[@]}"; do
+        if [[ -f "${c}" ]]; then
+            echo "${c}"
+            return 0
+        fi
+    done
+    local found=""
+    found="$(find "${home_dir}" /opt /usr/local /shared -maxdepth 6 -type f -path '*/etc/profile.d/conda.sh' 2>/dev/null | head -n 1 || true)"
+    if [[ -n "${found}" && -f "${found}" ]]; then
+        echo "${found}"
+        return 0
+    fi
+    return 1
+}
+
+CONDA_SH="$(resolve_conda_sh || true)"
+if [[ -z "${CONDA_SH}" || ! -f "${CONDA_SH}" ]]; then
+    echo "ERROR: conda.sh not found. Set CONDA_SH explicitly." >&2
+    exit 1
+fi
+source "${CONDA_SH}"
+eval "$(conda shell.bash hook)"
+conda activate "${CONDA_ENV}"
+
+cd "${PROJECT_DIR}"
+export PYTHONPATH="${PROJECT_DIR}:${PYTHONPATH:-}"
+
+JOB_TAG="transformer_labs_s${SEED}_c${_CONDOR_CLUSTER_ID:-local}_p${_CONDOR_PROCNO:-0}"
+
+echo "========================================================"
+echo "  Labs-only Transformer run  |  ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "  Conda env : ${CONDA_ENV}"
+echo "  EHR root  : ${EHR_ROOT}"
+echo "  Cache dir : ${CACHE_DIR}"
+echo "  Output dir: ${OUTPUT_DIR}"
+echo "  Seed      : ${SEED}"
+echo "  Dev mode  : ${DEV_MODE}"
+echo "========================================================"
+
+if ! python -c "import pyhealth" >/dev/null 2>&1; then
+    echo "ERROR: pyhealth is not importable. Run: bash condor_setup.sh" >&2
+    exit 1
+fi
+
+COMMON=(
+    --ehr-root "${EHR_ROOT}"
+    --cache-dir "${CACHE_DIR}"
+    --task labs
+    --model transformer
+    --embedding-dim "${EMBEDDING_DIM}"
+    --heads "${HEADS}"
+    --num-layers "${NUM_LAYERS}"
+    --dropout "${DROPOUT}"
+    --epochs "${EPOCHS}"
+    --batch-size "${BATCH_SIZE}"
+    --lr "${LR}"
+    --weight-decay "${WEIGHT_DECAY}"
+    --patience "${PATIENCE}"
+    --num-workers "${NUM_WORKERS}"
+    --pos-weight "${POS_WEIGHT}"
+    --seed "${SEED}"
+    --output-dir "${OUTPUT_DIR}"
+)
+
+if [[ "${DEV_MODE}" == "1" ]]; then
+    COMMON+=(--dev)
+fi
+
+if [[ "${USE_WANDB}" == "1" ]]; then
+    COMMON+=(--wandb --wandb-project "${WANDB_PROJECT}")
+    if [[ -n "${WANDB_RUN_NAME}" ]]; then
+        COMMON+=(--wandb-run-name "${WANDB_RUN_NAME}")
+    fi
+fi
+
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py "${COMMON[@]}"
+
+echo "========================================================"
+echo "  Completed ${JOB_TAG}"
+echo "  $(date -u '+%Y-%m-%d %H:%M:%S UTC')"
+echo "========================================================"
diff --git a/scripts/will/lambda-labs/labs/tmux_run_labs_bottleneck_transformer_variant.py b/scripts/will/lambda-labs/labs/tmux_run_labs_bottleneck_transformer_variant.py
new file mode 100644
index 000000000..0cc06f965
--- /dev/null
+++ b/scripts/will/lambda-labs/labs/tmux_run_labs_bottleneck_transformer_variant.py
@@ -0,0 +1,123 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+bottlenecks_n = 4
+fusion_startidx = 1
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"bottleneck_transformer_labs_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"bottleneck_transformer_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model bottleneck_transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--bottlenecks-n {bottlenecks_n}",
+    f"--fusion-startidx {fusion_startidx}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/lambda-labs/labs/tmux_run_labs_ehrmamba_variant.py b/scripts/will/lambda-labs/labs/tmux_run_labs_ehrmamba_variant.py
new file mode 100644
index 000000000..984db14fa
--- /dev/null
+++ b/scripts/will/lambda-labs/labs/tmux_run_labs_ehrmamba_variant.py
@@ -0,0 +1,121 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+num_layers = 2
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"ehrmamba_labs_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"ehrmamba_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model ehrmamba",
+    f"--embedding-dim {embedding_dim}",
+    f"--num-layers {num_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/lambda-labs/labs/tmux_run_labs_jambaehr_variant.py b/scripts/will/lambda-labs/labs/tmux_run_labs_jambaehr_variant.py
new file mode 100644
index 000000000..4c37c991f
--- /dev/null
+++ b/scripts/will/lambda-labs/labs/tmux_run_labs_jambaehr_variant.py
@@ -0,0 +1,125 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+jamba_transformer_layers = 2
+jamba_mamba_layers = 6
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-4
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"jambaehr_labs_s{seed}_batchsize{batch_size}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"jambaehr_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model jambaehr",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--jamba-transformer-layers {jamba_transformer_layers}",
+    f"--jamba-mamba-layers {jamba_mamba_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/lambda-labs/labs/tmux_run_labs_rnn_variant.py b/scripts/will/lambda-labs/labs/tmux_run_labs_rnn_variant.py
new file mode 100644
index 000000000..ddfbd219e
--- /dev/null
+++ b/scripts/will/lambda-labs/labs/tmux_run_labs_rnn_variant.py
@@ -0,0 +1,121 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"rnn_labs_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"rnn_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model rnn",
+    f"--embedding-dim {embedding_dim}",
+    f"--hidden-dim {hidden_dim}",
+    f"--rnn-type {rnn_type}",
+    f"--rnn-layers {rnn_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/lambda-labs/labs/tmux_run_labs_transformer_variant.py b/scripts/will/lambda-labs/labs/tmux_run_labs_transformer_variant.py
new file mode 100644
index 000000000..99afa77c8
--- /dev/null
+++ b/scripts/will/lambda-labs/labs/tmux_run_labs_transformer_variant.py
@@ -0,0 +1,119 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"transformer_labs_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"transformer_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_bottleneck_transformer_variant.py b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_bottleneck_transformer_variant.py
new file mode 100644
index 000000000..e3c36b43b
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_bottleneck_transformer_variant.py
@@ -0,0 +1,141 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+bottlenecks_n = 4
+fusion_startidx = 1
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"bottleneck_transformer_labs_notes_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"bottleneck_transformer_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model bottleneck_transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--bottlenecks-n {bottlenecks_n}",
+    f"--fusion-startidx {fusion_startidx}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_ehrmamba_variant.py b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_ehrmamba_variant.py
new file mode 100644
index 000000000..613782601
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_ehrmamba_variant.py
@@ -0,0 +1,139 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+num_layers = 2
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"ehrmamba_labs_notes_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"ehrmamba_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model ehrmamba",
+    f"--embedding-dim {embedding_dim}",
+    f"--num-layers {num_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_jambaehr_variant.py b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_jambaehr_variant.py
new file mode 100644
index 000000000..d01541f0d
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_jambaehr_variant.py
@@ -0,0 +1,143 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+jamba_transformer_layers = 2
+jamba_mamba_layers = 6
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-4
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"jambaehr_labs_notes_s{seed}_batchsize{batch_size}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"jambaehr_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model jambaehr",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--jamba-transformer-layers {jamba_transformer_layers}",
+    f"--jamba-mamba-layers {jamba_mamba_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_rnn_variant.py b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_rnn_variant.py
new file mode 100644
index 000000000..10012c522
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_rnn_variant.py
@@ -0,0 +1,139 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"rnn_labs_notes_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"rnn_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model rnn",
+    f"--embedding-dim {embedding_dim}",
+    f"--hidden-dim {hidden_dim}",
+    f"--rnn-type {rnn_type}",
+    f"--rnn-layers {rnn_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_transformer_variant.py b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_transformer_variant.py
new file mode 100644
index 000000000..b209e4857
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes/tmux_run_labs_notes_transformer_variant.py
@@ -0,0 +1,137 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"transformer_labs_notes_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"transformer_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_bottleneck_transformer_variant.py b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_bottleneck_transformer_variant.py
new file mode 100644
index 000000000..a7817267b
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_bottleneck_transformer_variant.py
@@ -0,0 +1,145 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cxr_root = "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0"
+cxr_variant = "default"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+bottlenecks_n = 4
+fusion_startidx = 1
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = True
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes-cxr"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"bottleneck_transformer_labs_notes_cxr_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"bottleneck_transformer_labs_notes_cxr_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cxr-root {cxr_root}",
+    f"--cxr-variant {cxr_variant}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs_cxr{' --dev' if dev else ''}",
+    "--model bottleneck_transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--bottlenecks-n {bottlenecks_n}",
+    f"--fusion-startidx {fusion_startidx}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsCXRMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_ehrmamba_variant.py b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_ehrmamba_variant.py
new file mode 100644
index 000000000..da321027f
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_ehrmamba_variant.py
@@ -0,0 +1,143 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cxr_root = "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0"
+cxr_variant = "default"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+num_layers = 2
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes-cxr"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"ehrmamba_labs_notes_cxr_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"ehrmamba_labs_notes_cxr_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cxr-root {cxr_root}",
+    f"--cxr-variant {cxr_variant}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs_cxr{' --dev' if dev else ''}",
+    "--model ehrmamba",
+    f"--embedding-dim {embedding_dim}",
+    f"--num-layers {num_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsCXRMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_jambaehr_variant.py b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_jambaehr_variant.py
new file mode 100644
index 000000000..70b99d4cd
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_jambaehr_variant.py
@@ -0,0 +1,147 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cxr_root = "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0"
+cxr_variant = "default"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+jamba_transformer_layers = 2
+jamba_mamba_layers = 6
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-4
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes-cxr"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"jambaehr_labs_notes_cxr_s{seed}_batchsize{batch_size}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"jambaehr_labs_notes_cxr_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cxr-root {cxr_root}",
+    f"--cxr-variant {cxr_variant}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs_cxr{' --dev' if dev else ''}",
+    "--model jambaehr",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--jamba-transformer-layers {jamba_transformer_layers}",
+    f"--jamba-mamba-layers {jamba_mamba_layers}",
+    f"--mamba-state-size {mamba_state_size}",
+    f"--mamba-conv-kernel {mamba_conv_kernel}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsCXRMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py
new file mode 100644
index 000000000..159a195bb
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_rnn_variant.py
@@ -0,0 +1,143 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cxr_root = "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0"
+cxr_variant = "default"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes-cxr"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"rnn_labs_notes_cxr_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"rnn_labs_notes_cxr_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cxr-root {cxr_root}",
+    f"--cxr-variant {cxr_variant}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs_cxr{' --dev' if dev else ''}",
+    "--model rnn",
+    f"--embedding-dim {embedding_dim}",
+    f"--hidden-dim {hidden_dim}",
+    f"--rnn-type {rnn_type}",
+    f"--rnn-layers {rnn_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsCXRMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_transformer_variant.py b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_transformer_variant.py
new file mode 100644
index 000000000..2add49db1
--- /dev/null
+++ b/scripts/will/lambda-labs/labs_notes_cxr/tmux_run_labs_notes_cxr_transformer_variant.py
@@ -0,0 +1,141 @@
+project_dir = "PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/home/ubuntu/mimiciv-data/ehr"
+note_root = "/home/ubuntu/mimiciv-data"
+cxr_root = "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0"
+cxr_variant = "default"
+cache_dir = "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr"
+logs_dir = "/home/ubuntu/logs"
+output_dir = "/home/ubuntu/output"
+embedding_dim = 128
+heads = 4
+num_layers = 2
+dropout = 0.1
+use_amp = True
+amp_dtype = "bf16"
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+pos_weight = 1
+num_workers = 4
+freeze_encoder = True
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes-cxr"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "0"
+session_name = f"transformer_labs_notes_cxr_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"transformer_labs_notes_cxr_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cxr-root {cxr_root}",
+    f"--cxr-variant {cxr_variant}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs_cxr{' --dev' if dev else ''}",
+    "--model transformer",
+    f"--embedding-dim {embedding_dim}",
+    f"--heads {heads}",
+    f"--num-layers {num_layers}",
+    f"--dropout {dropout}",
+]
+if use_amp:
+    flags.append("--use-amp")
+    flags.append(f"--amp-dtype {amp_dtype}")
+flags += [
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--pos-weight {pos_weight}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsCXRMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/read_cache.ipynb b/scripts/will/read_cache.ipynb
new file mode 100644
index 000000000..9af05b71b
--- /dev/null
+++ b/scripts/will/read_cache.ipynb
@@ -0,0 +1,94 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "711abd77-4eb5-4589-8200-e26de83fc129",
+   "metadata": {},
+   "source": [
+    "# Read Cache"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6354ebb1-bb74-4ec2-ad96-72e9ec869e36",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import pandas as pd"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "b03418f7-7a3f-44a9-9ff9-57cef85d3017",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "The current working directory is: /home/wp14\n"
+     ]
+    }
+   ],
+   "source": [
+    "os.chdir(\"../../../\")\n",
+    "print(f\"The current working directory is: {os.getcwd()}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "458c9daa-f889-4537-b699-a750903644e8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df = pd.read_parquet(\"pyhealth_cache_labs/012dadb7-77c0-575b-aef5-f1d7fee6a314/global_event_df.parquet\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "85be3049-e30d-4461-b645-730a79eaa770",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(118975491, 33)"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "df.shape"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/scripts/will/slurm/labs_only/slurm_run_labs_only_variant.py b/scripts/will/slurm/labs_only/slurm_run_labs_only_variant.py
new file mode 100644
index 000000000..6e13a38ea
--- /dev/null
+++ b/scripts/will/slurm/labs_only/slurm_run_labs_only_variant.py
@@ -0,0 +1,116 @@
+project_dir = "/u/wp14/PyHealth"
+seed = 12
+account = "jimeng-ic"
+partition = "eng-research-gpu"
+time = "24:00:00"
+mem = "64G"
+gres = "gpu:A10:1"
+conda_env = "pyhealth2"
+ehr_root = "/projects/illinois/eng/cs/jimeng/physionet.org/files/mimiciv/2.2"
+cache_dir = "/u/wp14/pyhealth_cache"
+output_dir = "output/rnn_labs"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+epochs = 5
+batch_size = 32
+lr = 1e-3 
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+dev = False
+
+# ── Step 0: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0 (optional): Clean logs and cache
+
+rm -rf {project_dir}/logs/*
+rm -rf {cache_dir}/*
+""")
+
+# ── Step 1: Reserve resources ─────────────────────────────────────────────────
+if dev:
+    print(f"""
+### STEP 1: Reserve an interactive node
+
+srun \\
+    --account={account} \\
+    --partition={partition} \\
+    --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\
+    --mem={mem} --gres={gres} --time={time} \\
+    --pty bash
+""")
+else:
+    print(f"""
+### STEP 1: Submit batch job (includes step 2 automatically)
+
+mkdir -p {project_dir}/logs && sbatch \\
+    --job-name=rnn_labs_s{seed} \\
+    --account={account} \\
+    --partition={partition} \\
+    --nodes=1 --ntasks=1 --cpus-per-task={num_workers} \\
+    --mem={mem} --gres={gres} --time={time} \\
+    --output={project_dir}/logs/rnn_labs_s{seed}_%j.out \\
+    --error={project_dir}/logs/rnn_labs_s{seed}_%j.err \\
+    --wrap='
+        module load miniconda3/24.9.2 &&
+        eval "$(conda shell.bash hook)" &&
+        conda activate {conda_env} &&
+        cd {project_dir} &&
+        export PYTHONPATH={project_dir}:$PYTHONPATH &&
+        python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+            --ehr-root {ehr_root} \\
+            --cache-dir {cache_dir} \\
+            --task labs \\
+            --model rnn \\
+            --embedding-dim {embedding_dim} \\
+            --hidden-dim {hidden_dim} \\
+            --rnn-type {rnn_type} \\
+            --rnn-layers {rnn_layers} \\
+            --dropout {dropout} \\
+            --epochs {epochs} \\
+            --batch-size {batch_size} \\
+            --lr {lr} \\
+            --weight-decay {weight_decay} \\
+            --patience {patience} \\
+            --num-workers {num_workers} \\
+            --seed {seed} \\
+            --output-dir {output_dir}
+    '
+""")
+
+# ── Step 2: Run (only needed for dev/interactive) ─────────────────────────────
+print("\n" + "=" * 60 + "\n")
+if dev:
+    print(f"""
+### STEP 2: Once on the node, run
+
+module load miniconda3/24.9.2 &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+mkdir -p {project_dir}/logs/dev &&
+python {project_dir}/examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    --ehr-root {ehr_root} \\
+    --cache-dir {cache_dir} \\
+    --task labs \\
+    --model rnn \\
+    --embedding-dim {embedding_dim} \\
+    --hidden-dim {hidden_dim} \\
+    --rnn-type {rnn_type} \\
+    --rnn-layers {rnn_layers} \\
+    --dropout {dropout} \\
+    --epochs {epochs} \\
+    --batch-size {batch_size} \\
+    --lr {lr} \\
+    --weight-decay {weight_decay} \\
+    --patience {patience} \\
+    --num-workers {num_workers} \\
+    --seed {seed} \\
+    --output-dir {output_dir} --dev \\
+    > >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.out) \\
+    2> >(tee {project_dir}/logs/dev/rnn_labs_s{seed}.err >&2)
+""")
diff --git a/scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py b/scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py
new file mode 100644
index 000000000..fe8704565
--- /dev/null
+++ b/scripts/will/sunlab/labs_notes/tmux_run_labs_notes_rnn_variant.py
@@ -0,0 +1,130 @@
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+cache_dir = "/shared/eng/wp14/pyhealth_cache_labs_notes"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+freeze_encoder = False
+include_vitals = False
+dev = False
+use_old_cache = False
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-notes"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "2"
+session_name = f"rnn_labs_notes_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"rnn_labs_notes_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--note-root {note_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task notes_labs{' --dev' if dev else ''}",
+    "--model rnn",
+    f"--embedding-dim {embedding_dim}",
+    f"--hidden-dim {hidden_dim}",
+    f"--rnn-type {rnn_type}",
+    f"--rnn-layers {rnn_layers}",
+    f"--dropout {dropout}",
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--num-workers {num_workers}",
+]
+if freeze_encoder:
+    flags.append("--freeze-encoder")
+if include_vitals:
+    flags.append("--include-vitals")
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Confirm which note sections were used for filtering ─────────────
+print(f"""
+### STEP 4: Confirm which discharge-note AND radiology-note sections
+### NotesLabsMIMIC4 filtered to (e.g. discharge: "chief complaint";
+### radiology: "indication", "impression"). Logged once at INFO level by
+### pyhealth.tasks.multimodal_mimic4 when the task is constructed.
+
+grep "filtering discharge notes to sections" {log_dir}/{log_tag}.out
+""")
diff --git a/scripts/will/sunlab/labs_only/tmux_run_labs_only_bottleneck_transformer_variant.py b/scripts/will/sunlab/labs_only/tmux_run_labs_only_bottleneck_transformer_variant.py
new file mode 100644
index 000000000..38d9afeec
--- /dev/null
+++ b/scripts/will/sunlab/labs_only/tmux_run_labs_only_bottleneck_transformer_variant.py
@@ -0,0 +1,140 @@
+from datetime import datetime
+
+date_str = datetime.now().strftime("%Y%m%d")
+
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+task = "labs"
+cache_dir = f"/home/wp14/pyhealth_cache_{task}"
+embedding_dim = 64
+hidden_dim = 64
+heads = 4
+num_layers = 2
+bottlenecks_n = 4
+fusion_startidx = 1
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+pos_weight = 1.0
+max_grad_norm = 0.5
+dev = False
+cuda_visible_devices = "4"
+use_old_cache = True
+session_name = f"bottleneck_transformer_{task}_s{seed}"
+wandb_project = "pyhealth-multimodal-labs-only"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"bottleneck_transformer_{task}_s{seed}_{date_str}"
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+export WANDB_PROJECT={wandb_project} &&
+export WANDB_NAME={log_tag} &&
+{{ echo "[params] seed={seed} task={task} cache_dir={cache_dir} embedding_dim={embedding_dim} hidden_dim={hidden_dim} heads={heads} num_layers={num_layers} bottlenecks_n={bottlenecks_n} fusion_startidx={fusion_startidx} dropout={dropout} epochs={epochs} batch_size={batch_size} lr={lr} weight_decay={weight_decay} patience={patience} num_workers={num_workers} pos_weight={pos_weight} max_grad_norm={max_grad_norm} dev={dev}" | tee {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+{{ echo "[start] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    --ehr-root {ehr_root} \\
+    --cache-dir {cache_dir} \\
+    --task {task}{' --dev' if dev else ''} \\
+    --model bottleneck_transformer \\
+    --embedding-dim {embedding_dim} \\
+    --hidden-dim {hidden_dim} \\
+    --heads {heads} \\
+    --num-layers {num_layers} \\
+    --bottlenecks-n {bottlenecks_n} \\
+    --fusion-startidx {fusion_startidx} \\
+    --dropout {dropout} \\
+    --epochs {epochs} \\
+    --batch-size {batch_size} \\
+    --lr {lr} \\
+    --weight-decay {weight_decay} \\
+    --patience {patience} \\
+    --num-workers {num_workers} \\
+    --pos-weight {pos_weight} \\
+    --max-grad-norm {max_grad_norm} \\
+    --seed {seed} \\
+    --output-dir {output_dir} \\
+    > >(tee -a {log_dir}/{log_tag}.out) \\
+    2> >(tee -a {log_dir}/{log_tag}.err >&2);
+{{ echo "[end] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }}
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### wandb: run should appear at (look for "wandb: 🚀 View run" in the log)
+
+https://wandb.ai/<entity>/{wandb_project}
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Compute token stats (optional) ───────────────────────────────────
+token_stats_tag = f"token_stats_notes_labs_s{seed}_{date_str}"
+
+print(f"""
+### STEP 4 (optional): Compute token stats (run in a separate shell, no GPU needed)
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+python scripts/compute_token_stats.py \\
+    --ehr-root {ehr_root} \\
+    --note-root {note_root} \\
+    --cache-dir {cache_dir} \\
+    --task notes_labs{' --dev' if dev else ''} \\
+    > >(tee {log_dir}/{token_stats_tag}.out) \\
+    2> >(tee {log_dir}/{token_stats_tag}.err >&2)
+""")
diff --git a/scripts/will/sunlab/labs_only/tmux_run_labs_only_ehrmamba_variant.py b/scripts/will/sunlab/labs_only/tmux_run_labs_only_ehrmamba_variant.py
new file mode 100644
index 000000000..2ca0ed65a
--- /dev/null
+++ b/scripts/will/sunlab/labs_only/tmux_run_labs_only_ehrmamba_variant.py
@@ -0,0 +1,134 @@
+from datetime import datetime
+
+date_str = datetime.now().strftime("%Y%m%d")
+
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+task = "labs"
+cache_dir = f"/home/wp14/pyhealth_cache_{task}"
+embedding_dim = 64
+num_layers = 2
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+pos_weight = 1.0
+dev = False
+cuda_visible_devices = "4"
+use_old_cache = True
+session_name = f"ehrmamba_{task}_s{seed}"
+wandb_project = "pyhealth-multimodal-labs-only"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"ehrmamba_{task}_s{seed}_{date_str}"
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+export WANDB_PROJECT={wandb_project} &&
+export WANDB_NAME={log_tag} &&
+{{ echo "[params] seed={seed} task={task} cache_dir={cache_dir} embedding_dim={embedding_dim} num_layers={num_layers} mamba_state_size={mamba_state_size} mamba_conv_kernel={mamba_conv_kernel} dropout={dropout} epochs={epochs} batch_size={batch_size} lr={lr} weight_decay={weight_decay} patience={patience} num_workers={num_workers} pos_weight={pos_weight} dev={dev}" | tee {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+{{ echo "[start] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    --ehr-root {ehr_root} \\
+    --cache-dir {cache_dir} \\
+    --task {task}{' --dev' if dev else ''} \\
+    --model ehrmamba \\
+    --embedding-dim {embedding_dim} \\
+    --num-layers {num_layers} \\
+    --mamba-state-size {mamba_state_size} \\
+    --mamba-conv-kernel {mamba_conv_kernel} \\
+    --dropout {dropout} \\
+    --epochs {epochs} \\
+    --batch-size {batch_size} \\
+    --lr {lr} \\
+    --weight-decay {weight_decay} \\
+    --patience {patience} \\
+    --num-workers {num_workers} \\
+    --pos-weight {pos_weight} \\
+    --seed {seed} \\
+    --output-dir {output_dir} \\
+    > >(tee -a {log_dir}/{log_tag}.out) \\
+    2> >(tee -a {log_dir}/{log_tag}.err >&2);
+{{ echo "[end] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }}
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### wandb: run should appear at (look for "wandb: 🚀 View run" in the log)
+
+https://wandb.ai/<entity>/{wandb_project}
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Compute token stats (optional) ───────────────────────────────────
+token_stats_tag = f"token_stats_notes_labs_s{seed}_{date_str}"
+
+print(f"""
+### STEP 4 (optional): Compute token stats (run in a separate shell, no GPU needed)
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+python scripts/compute_token_stats.py \\
+    --ehr-root {ehr_root} \\
+    --note-root {note_root} \\
+    --cache-dir {cache_dir} \\
+    --task notes_labs{' --dev' if dev else ''} \\
+    > >(tee {log_dir}/{token_stats_tag}.out) \\
+    2> >(tee {log_dir}/{token_stats_tag}.err >&2)
+""")
diff --git a/scripts/will/sunlab/labs_only/tmux_run_labs_only_jambaehr_variant.py b/scripts/will/sunlab/labs_only/tmux_run_labs_only_jambaehr_variant.py
new file mode 100644
index 000000000..c1b3ec6f1
--- /dev/null
+++ b/scripts/will/sunlab/labs_only/tmux_run_labs_only_jambaehr_variant.py
@@ -0,0 +1,138 @@
+from datetime import datetime
+
+date_str = datetime.now().strftime("%Y%m%d")
+
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+task = "labs"
+cache_dir = f"/home/wp14/pyhealth_cache_{task}"
+embedding_dim = 64
+heads = 4
+jamba_transformer_layers = 2
+jamba_mamba_layers = 6
+mamba_state_size = 16
+mamba_conv_kernel = 4
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+pos_weight = 1.0
+dev = False
+cuda_visible_devices = "4"
+use_old_cache = True
+session_name = f"jambaehr_{task}_s{seed}"
+wandb_project = "pyhealth-multimodal-labs-only"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"jambaehr_{task}_s{seed}_{date_str}"
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+export WANDB_PROJECT={wandb_project} &&
+export WANDB_NAME={log_tag} &&
+{{ echo "[params] seed={seed} task={task} cache_dir={cache_dir} embedding_dim={embedding_dim} heads={heads} jamba_transformer_layers={jamba_transformer_layers} jamba_mamba_layers={jamba_mamba_layers} mamba_state_size={mamba_state_size} mamba_conv_kernel={mamba_conv_kernel} dropout={dropout} epochs={epochs} batch_size={batch_size} lr={lr} weight_decay={weight_decay} patience={patience} num_workers={num_workers} pos_weight={pos_weight} dev={dev}" | tee {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+{{ echo "[start] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    --ehr-root {ehr_root} \\
+    --cache-dir {cache_dir} \\
+    --task {task}{' --dev' if dev else ''} \\
+    --model jambaehr \\
+    --embedding-dim {embedding_dim} \\
+    --heads {heads} \\
+    --jamba-transformer-layers {jamba_transformer_layers} \\
+    --jamba-mamba-layers {jamba_mamba_layers} \\
+    --mamba-state-size {mamba_state_size} \\
+    --mamba-conv-kernel {mamba_conv_kernel} \\
+    --dropout {dropout} \\
+    --epochs {epochs} \\
+    --batch-size {batch_size} \\
+    --lr {lr} \\
+    --weight-decay {weight_decay} \\
+    --patience {patience} \\
+    --num-workers {num_workers} \\
+    --pos-weight {pos_weight} \\
+    --seed {seed} \\
+    --output-dir {output_dir} \\
+    > >(tee -a {log_dir}/{log_tag}.out) \\
+    2> >(tee -a {log_dir}/{log_tag}.err >&2);
+{{ echo "[end] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }}
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### wandb: run should appear at (look for "wandb: 🚀 View run" in the log)
+
+https://wandb.ai/<entity>/{wandb_project}
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Compute token stats (optional) ───────────────────────────────────
+token_stats_tag = f"token_stats_notes_labs_s{seed}_{date_str}"
+
+print(f"""
+### STEP 4 (optional): Compute token stats (run in a separate shell, no GPU needed)
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+python scripts/compute_token_stats.py \\
+    --ehr-root {ehr_root} \\
+    --note-root {note_root} \\
+    --cache-dir {cache_dir} \\
+    --task notes_labs{' --dev' if dev else ''} \\
+    > >(tee {log_dir}/{token_stats_tag}.out) \\
+    2> >(tee {log_dir}/{token_stats_tag}.err >&2)
+""")
diff --git a/scripts/will/sunlab/labs_only/tmux_run_labs_only_rnn_variant.py b/scripts/will/sunlab/labs_only/tmux_run_labs_only_rnn_variant.py
new file mode 100644
index 000000000..abbbbef2e
--- /dev/null
+++ b/scripts/will/sunlab/labs_only/tmux_run_labs_only_rnn_variant.py
@@ -0,0 +1,112 @@
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+cache_dir = "/shared/eng/wp14/pyhealth_cache_labs"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+embedding_dim = 128
+hidden_dim = 128
+rnn_type = "GRU"
+rnn_layers = 2
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+dev = False
+use_old_cache = True
+use_wandb = True
+wandb_project = "pyhealth-multimodal-labs-only"
+wandb_run_name = None  # defaults to "{model}_seed{seed}" if unset
+cuda_visible_devices = "4"
+session_name = f"rnn_labs_s{seed}"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"rnn_labs_s{seed}"
+
+flags = [
+    f"--ehr-root {ehr_root}",
+    f"--cache-dir {cache_dir}",
+    f"--task labs{' --dev' if dev else ''}",
+    "--model rnn",
+    f"--embedding-dim {embedding_dim}",
+    f"--hidden-dim {hidden_dim}",
+    f"--rnn-type {rnn_type}",
+    f"--rnn-layers {rnn_layers}",
+    f"--dropout {dropout}",
+    f"--epochs {epochs}",
+    f"--batch-size {batch_size}",
+    f"--lr {lr}",
+    f"--weight-decay {weight_decay}",
+    f"--patience {patience}",
+    f"--num-workers {num_workers}",
+]
+if use_wandb:
+    flags.append("--wandb")
+    flags.append(f"--wandb-project {wandb_project}")
+    if wandb_run_name:
+        flags.append(f"--wandb-run-name {wandb_run_name}")
+flags += [
+    f"--seed {seed}",
+    f"--output-dir {output_dir}",
+]
+flags_block = " \\\n    ".join(flags)
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    {flags_block} \\
+    > >(tee {log_dir}/{log_tag}.out) \\
+    2> >(tee {log_dir}/{log_tag}.err >&2)
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
diff --git a/scripts/will/sunlab/labs_only/tmux_run_labs_only_transformer_variant.py b/scripts/will/sunlab/labs_only/tmux_run_labs_only_transformer_variant.py
new file mode 100644
index 000000000..153189d15
--- /dev/null
+++ b/scripts/will/sunlab/labs_only/tmux_run_labs_only_transformer_variant.py
@@ -0,0 +1,134 @@
+from datetime import datetime
+
+date_str = datetime.now().strftime("%Y%m%d")
+
+project_dir = "/home/wp14/PyHealth"
+seed = 12
+conda_env = "pyhealth2"
+ehr_root = "/shared/rsaas/physionet.org/files/mimiciv/2.2"
+note_root = "/shared/rsaas/physionet.org/files/mimic-note"
+logs_dir = "/home/wp14/logs"
+output_dir = "/home/wp14/output"
+task = "labs"
+cache_dir = f"/home/wp14/pyhealth_cache_{task}"
+embedding_dim = 64
+hidden_dim = 64
+heads = 4
+num_layers = 2
+dropout = 0.1
+epochs = 50
+batch_size = 32
+lr = 1e-3
+weight_decay = 1e-5
+patience = 5
+num_workers = 4
+pos_weight = 1.0
+dev = False
+cuda_visible_devices = "4"
+use_old_cache = False
+session_name = f"transformer_{task}_s{seed}"
+wandb_project = "pyhealth-multimodal-labs-only"
+
+# ── Step 0a: Check which CUDA GPU is available ───────────────────────────────
+print(f"""
+### STEP 0a: Check which CUDA GPU is available (pick an idle index for CUDA_VISIBLE_DEVICES)
+
+nvidia-smi --query-gpu=index,utilization.gpu,memory.used,memory.free --format=csv
+""")
+
+# ── Step 0b: Clean logs and cache ────────────────────────────────────────────
+print(f"""
+### STEP 0b (optional): Clean logs and cache
+
+rm -rf {logs_dir}/*
+{'# cache preserved (use_old_cache=True)' if use_old_cache else f'rm -rf {cache_dir}/*'}
+rm -rf {output_dir}/*
+""")
+
+# ── Step 1: Start a tmux session ────────────────────────────────────────────
+print(f"""
+### STEP 1: Start a tmux session (attached)
+
+tmux new-session -s {session_name}
+""")
+
+# ── Step 2: Paste this into the session to run training ─────────────────────
+print("\n" + "=" * 60 + "\n")
+
+log_dir = logs_dir
+log_tag = f"transformer_{task}_s{seed}_{date_str}"
+
+print(f"""
+### STEP 2: Paste this into the tmux session
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+export PYTHONPATH={project_dir}:$PYTHONPATH &&
+export CUDA_VISIBLE_DEVICES={cuda_visible_devices} &&
+export WANDB_PROJECT={wandb_project} &&
+export WANDB_NAME={log_tag} &&
+{{ echo "[params] seed={seed} task={task} cache_dir={cache_dir} embedding_dim={embedding_dim} hidden_dim={hidden_dim} heads={heads} num_layers={num_layers} dropout={dropout} epochs={epochs} batch_size={batch_size} lr={lr} weight_decay={weight_decay} patience={patience} num_workers={num_workers} pos_weight={pos_weight} dev={dev}" | tee {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+{{ echo "[start] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }} &&
+python examples/mortality_prediction/unified_embedding_e2e_mimic4.py \\
+    --ehr-root {ehr_root} \\
+    --cache-dir {cache_dir} \\
+    --task {task}{' --dev' if dev else ''} \\
+    --model transformer \\
+    --embedding-dim {embedding_dim} \\
+    --hidden-dim {hidden_dim} \\
+    --heads {heads} \\
+    --num-layers {num_layers} \\
+    --dropout {dropout} \\
+    --epochs {epochs} \\
+    --batch-size {batch_size} \\
+    --lr {lr} \\
+    --weight-decay {weight_decay} \\
+    --patience {patience} \\
+    --num-workers {num_workers} \\
+    --pos-weight {pos_weight} \\
+    --seed {seed} \\
+    --output-dir {output_dir} \\
+    > >(tee -a {log_dir}/{log_tag}.out) \\
+    2> >(tee -a {log_dir}/{log_tag}.err >&2);
+{{ echo "[end] $(TZ='America/Los_Angeles' date '+%Y-%m-%d %H:%M:%S %Z')" | tee -a {log_dir}/{log_tag}.out | tee -a {log_dir}/{log_tag}.err > /dev/null; }}
+""")
+
+# ── Step 3: Detach / reattach / monitor ──────────────────────────────────────
+print(f"""
+### STEP 3: Detach without killing (Ctrl+b d), then reattach later with
+
+tmux attach -t {session_name}
+
+### To check on it later without attaching:
+
+tail -f {log_dir}/{log_tag}.out
+
+### wandb: run should appear at (look for "wandb: 🚀 View run" in the log)
+
+https://wandb.ai/<entity>/{wandb_project}
+
+### To kill the session when done:
+
+tmux kill-session -t {session_name}
+""")
+
+# ── Step 4: Compute token stats (optional) ───────────────────────────────────
+token_stats_tag = f"token_stats_notes_labs_s{seed}_{date_str}"
+
+print(f"""
+### STEP 4 (optional): Compute token stats (run in a separate shell, no GPU needed)
+
+mkdir -p {log_dir} &&
+eval "$(conda shell.bash hook)" &&
+conda activate {conda_env} &&
+cd {project_dir} &&
+python scripts/compute_token_stats.py \\
+    --ehr-root {ehr_root} \\
+    --note-root {note_root} \\
+    --cache-dir {cache_dir} \\
+    --task notes_labs{' --dev' if dev else ''} \\
+    > >(tee {log_dir}/{token_stats_tag}.out) \\
+    2> >(tee {log_dir}/{token_stats_tag}.err >&2)
+""")
diff --git a/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv
index 804e2ae40..21511f015 100644
--- a/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv
+++ b/test-resources/core/mimic4demo/hosp/diagnoses_icd.csv
@@ -19,6 +19,7 @@ subject_id,hadm_id,seq_num,icd_code,icd_version
 10003,20006,2,E1065,10
 10003,20006,3,N170,10
 10003,20006,4,I509,10
+10003,20006,5,R99,10
 10004,20007,1,K219,10
 10004,20007,2,I10,10
 10005,20008,1,25001,9
diff --git a/test-resources/core/mimic4demo/note/discharge.csv b/test-resources/core/mimic4demo/note/discharge.csv
new file mode 100644
index 000000000..ddc46151f
--- /dev/null
+++ b/test-resources/core/mimic4demo/note/discharge.csv
@@ -0,0 +1,1869 @@
+"note_id","subject_id","hadm_id","note_type","note_seq","charttime","storetime","text"
+"d1","10001","19999","DS","1","2150-02-18 12:00:00","2150-02-18 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: SURGERY
+
+Allergies:
+morphine
+
+Attending: ___.
+
+Chief Complaint:
+abdominal pain
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+This patient is a ___ year old male who complains of right lower quadrant
+abdominal pain for 2 days. Pain worsens with movement. Patient denies fevers
+or chills. No relief with over-the-counter analgesics.
+
+Past Medical History:
+none
+
+Social History:
+___
+
+Family History:
+NC
+
+Physical Exam:
+Temp: 97.8 HR: 90 BP: 124/86 Resp: 14 O2Sat: 100
+Abdomen: right lower quadrant tenderness without rebound
+
+Pertinent Results:
+WBC-8.9 RBC-5.59 Hgb-12.5 Hct-42.0
+Glucose-99 UreaN-13 Creat-1.0 Na-137 K-4.0 Cl-103 HCO3-22
+
+Brief Hospital Course:
+Patient admitted for observation. Serial abdominal exams performed.
+Pain resolved with conservative management. Diet advanced without issue.
+Discharged home in stable condition.
+
+Medications on Admission:
+none
+
+Discharge Medications:
+none
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+abdominal pain
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Return to ER for worsening pain, fever >101.5F, vomiting, or new symptoms.
+
+Followup Instructions:
+Follow up with primary care physician in 1 week."
+"d2","10002","20000","DS","1","2151-01-02 12:00:00","2151-01-02 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+penicillin
+
+Attending: ___.
+
+Chief Complaint:
+chest pain
+
+Major Surgical or Invasive Procedure:
+cardiac catheterization
+
+History of Present Illness:
+___ year old female with hypertension presenting with substernal chest pain
+radiating to the left arm, onset 3 hours prior to admission. Associated
+diaphoresis and shortness of breath. EKG showed ST depression in V4-V6.
+
+Past Medical History:
+1. Hypertension
+2. Hyperlipidemia
+3. Type 2 Diabetes Mellitus
+
+Social History:
+Non-smoker. Occasional alcohol use. Lives with spouse.
+
+Family History:
+Father with MI at age 60. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.2 HR: 88 BP: 152/94 Resp: 18 O2Sat: 97
+Cardiovascular: Regular rate and rhythm. No murmurs.
+Chest: Clear to auscultation bilaterally.
+
+Pertinent Results:
+Troponin I: 2.4 (elevated)
+BNP: 180
+WBC-9.1 Hgb-11.8 Hct-36.2 Plt-220
+
+Brief Hospital Course:
+Patient admitted for NSTEMI. Heparin drip initiated. Cardiac cath revealed
+70% LAD stenosis; drug-eluting stent placed. Post-procedure course
+uncomplicated. Discharged on dual antiplatelet therapy.
+
+Medications on Admission:
+1. Lisinopril 10 mg daily
+2. Atorvastatin 40 mg nightly
+3. Metformin 1000 mg twice daily
+
+Discharge Medications:
+1. Aspirin 81 mg daily
+2. Clopidogrel 75 mg daily
+3. Lisinopril 10 mg daily
+4. Atorvastatin 80 mg nightly
+5. Metformin 1000 mg twice daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+NSTEMI, coronary artery disease
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Do not stop aspirin or clopidogrel without consulting your cardiologist.
+Call your doctor for chest pain, shortness of breath, or leg swelling.
+
+Followup Instructions:
+Follow up with cardiology in 1 week."
+"d3","10001","20001","DS","1","2150-03-18 12:00:00","2150-03-18 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+sulfa drugs
+
+Attending: ___.
+
+Chief Complaint:
+shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male with known COPD presenting with 3 days of worsening dyspnea
+and productive cough with yellow sputum. Patient reports increased inhaler use
+without improvement. Denies fever. Smoking history: 40 pack-years.
+
+Past Medical History:
+1. COPD, moderate severity
+2. Hypertension
+3. Gastroesophageal reflux disease
+
+Social History:
+Former smoker, quit 5 years ago. No alcohol or illicit drug use. Lives alone.
+
+Family History:
+Mother with asthma. No family history of lung cancer.
+
+Physical Exam:
+Temp: 98.6 HR: 102 BP: 138/82 Resp: 24 O2Sat: 88% on room air
+Chest: Diffuse expiratory wheezes. Prolonged expiratory phase.
+Accessory muscle use noted.
+
+Pertinent Results:
+ABG: pH 7.34 pCO2 52 pO2 58 HCO3 28
+WBC-11.2 (elevated) CXR: hyperinflation, no consolidation
+
+Brief Hospital Course:
+Admitted for COPD exacerbation. Started on systemic steroids, nebulized
+bronchodilators, and supplemental oxygen. Azithromycin added for possible
+infectious trigger. O2 sats improved to 94% on 2L NC. Discharged on
+steroid taper.
+
+Medications on Admission:
+1. Albuterol inhaler PRN
+2. Tiotropium 18 mcg daily
+3. Lisinopril 5 mg daily
+4. Omeprazole 20 mg daily
+
+Discharge Medications:
+1. Prednisone 40 mg daily x 3 days then taper per schedule
+2. Azithromycin 250 mg daily x 3 days (2 doses remaining)
+3. Albuterol inhaler q4h PRN
+4. Tiotropium 18 mcg daily
+5. Lisinopril 5 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+COPD exacerbation
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with mild exertional dyspnea.
+
+Discharge Instructions:
+Complete the full steroid taper. Use albuterol inhaler as needed.
+Return to ER for oxygen saturation below 90%, severe dyspnea, or confusion.
+
+Followup Instructions:
+Follow up with pulmonology in 2 weeks."
+"d4","10001","20002","DS","1","2150-06-25 12:00:00","2150-06-25 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+no known drug allergies
+
+Attending: ___.
+
+Chief Complaint:
+leg swelling and shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old female with history of congestive heart failure (EF 30%) presenting
+with 1 week of progressive bilateral lower extremity edema and orthopnea.
+Weight gain of 8 lbs over past week. Patient reports dietary indiscretion
+with increased sodium intake.
+
+Past Medical History:
+1. Congestive heart failure, EF 30%
+2. Atrial fibrillation
+3. Hypertension
+4. Chronic kidney disease stage 3
+
+Social History:
+Non-smoker. No alcohol. Lives with daughter.
+
+Family History:
+Father with heart failure. Sister with hypertension.
+
+Physical Exam:
+Temp: 98.0 HR: 96 (irregularly irregular) BP: 148/90 Resp: 20 O2Sat: 94%
+JVD present. Bilateral crackles at lung bases.
+2+ pitting edema to knees bilaterally.
+
+Pertinent Results:
+BNP: 1840 (markedly elevated)
+Creatinine 1.8 (baseline 1.5)
+CXR: pulmonary vascular congestion, bilateral pleural effusions
+
+Brief Hospital Course:
+Admitted for decompensated CHF. IV furosemide initiated with good diuretic
+response (negative 3L over 48 hours). Transitioned to oral diuretics.
+Weight at discharge 4 lbs below admission weight. Cardiology consulted.
+
+Medications on Admission:
+1. Furosemide 40 mg daily
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Medications:
+1. Furosemide 80 mg daily (increased)
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+Acute decompensated heart failure
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Weigh yourself daily. If weight increases by more than 2 lbs in a day or
+5 lbs in a week, call your doctor immediately. Restrict sodium to 2g/day.
+
+Followup Instructions:
+Follow up with cardiologist in 3 days."
+"d5","10002","20003","DS","1","2151-01-12 12:00:00","2151-01-12 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+amoxicillin
+
+Attending: ___.
+
+Chief Complaint:
+fever and productive cough
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male presenting with 5 days of productive cough with green sputum,
+fever to 101.8F, and pleuritic chest pain. Patient reports decreased appetite
+and fatigue. No sick contacts. Up to date on vaccinations.
+
+Past Medical History:
+1. Type 2 Diabetes Mellitus
+2. Hypertension
+
+Social History:
+Non-smoker. No alcohol or drug use. Works as a teacher.
+
+Family History:
+No family history of pulmonary disease.
+
+Physical Exam:
+Temp: 101.6 HR: 108 BP: 128/78 Resp: 22 O2Sat: 93% on room air
+Chest: Decreased breath sounds and dullness to percussion at right base.
+Egophony present at right lower lobe.
+
+Pertinent Results:
+WBC-14.8 (elevated) with left shift
+Procalcitonin: 2.1 (elevated)
+CXR: right lower lobe consolidation consistent with pneumonia
+
+Brief Hospital Course:
+Admitted for community-acquired pneumonia. Started on ceftriaxone and
+azithromycin. O2 sats improved on 4L NC. Transitioned to oral antibiotics
+on hospital day 3. Glucose levels monitored and insulin sliding scale used.
+
+Medications on Admission:
+1. Metformin 500 mg twice daily
+2. Lisinopril 10 mg daily
+3. Aspirin 81 mg daily
+
+Discharge Medications:
+1. Levofloxacin 750 mg daily x 5 days (3 doses remaining)
+2. Metformin 500 mg twice daily
+3. Lisinopril 10 mg daily
+4. Aspirin 81 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+Community-acquired pneumonia
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Complete the full antibiotic course. Return for worsening shortness of breath,
+persistent fever, or hemoptysis.
+
+Followup Instructions:
+Follow up with primary care in 1 week for repeat CXR."
+"d6","10002","20004","DS","1","2151-04-20 12:00:00","2151-04-20 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: NEUROLOGY
+
+Allergies:
+codeine
+
+Attending: ___.
+
+Chief Complaint:
+sudden onset left-sided weakness
+
+Major Surgical or Invasive Procedure:
+IV tPA administration
+
+History of Present Illness:
+___ year old female with hypertension and atrial fibrillation brought by EMS
+with sudden onset left arm and leg weakness and facial droop noted 2 hours
+prior to arrival. Last known well time established. NIHSS 12 on arrival.
+CT head showed no hemorrhage. IV tPA administered within window.
+
+Past Medical History:
+1. Atrial fibrillation
+2. Hypertension
+3. Hyperlipidemia
+
+Social History:
+Non-smoker. Rare alcohol use. Retired. Lives alone.
+
+Family History:
+Father with stroke at age 72. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.4 HR: 84 (irregular) BP: 168/96 Resp: 16 O2Sat: 98%
+Neuro: left facial droop, left arm 3/5 strength, left leg 4/5 strength.
+Dysarthria present.
+
+Pertinent Results:
+MRI brain: right MCA territory infarct
+Echo: no thrombus. EF 55%.
+INR: 1.1 (not anticoagulated at time of stroke)
+
+Brief Hospital Course:
+Admitted to stroke unit. tPA given without hemorrhagic complication.
+Deficits improved over 48 hours. Anticoagulation initiated with apixaban.
+Physical therapy and speech therapy consulted. Swallow evaluation passed.
+
+Medications on Admission:
+1. Metoprolol 25 mg twice daily
+2. Amlodipine 5 mg daily
+3. Atorvastatin 20 mg nightly
+
+Discharge Medications:
+1. Apixaban 5 mg twice daily (new)
+2. Aspirin 81 mg daily
+3. Metoprolol 25 mg twice daily
+4. Amlodipine 5 mg daily
+5. Atorvastatin 80 mg nightly (increased)
+
+Discharge Disposition:
+Rehabilitation facility
+
+Discharge Diagnosis:
+Acute ischemic stroke, right MCA territory
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Take apixaban as directed. Do not stop without consulting your neurologist.
+Call 911 immediately for any new weakness, vision changes, or speech difficulty.
+
+Followup Instructions:
+Follow up with neurology in 2 weeks."
+"d7","10003","20005","DS","1","2152-03-05 12:00:00","2152-03-05 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: SURGERY
+
+Allergies:
+morphine
+
+Attending: ___.
+
+Chief Complaint:
+abdominal pain
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+This patient is a ___ year old male who complains of right lower quadrant
+abdominal pain for 2 days. Pain worsens with movement. Patient denies fevers
+or chills. No relief with over-the-counter analgesics.
+
+Past Medical History:
+none
+
+Social History:
+___
+
+Family History:
+NC
+
+Physical Exam:
+Temp: 97.8 HR: 90 BP: 124/86 Resp: 14 O2Sat: 100
+Abdomen: right lower quadrant tenderness without rebound
+
+Pertinent Results:
+WBC-8.9 RBC-5.59 Hgb-12.5 Hct-42.0
+Glucose-99 UreaN-13 Creat-1.0 Na-137 K-4.0 Cl-103 HCO3-22
+
+Brief Hospital Course:
+Patient admitted for observation. Serial abdominal exams performed.
+Pain resolved with conservative management. Diet advanced without issue.
+Discharged home in stable condition.
+
+Medications on Admission:
+none
+
+Discharge Medications:
+none
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+abdominal pain
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Return to ER for worsening pain, fever >101.5F, vomiting, or new symptoms.
+
+Followup Instructions:
+Follow up with primary care physician in 1 week."
+"d8","10003","20006","DS","1","2152-08-15 12:00:00","2152-08-15 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+penicillin
+
+Attending: ___.
+
+Chief Complaint:
+chest pain
+
+Major Surgical or Invasive Procedure:
+cardiac catheterization
+
+History of Present Illness:
+___ year old female with hypertension presenting with substernal chest pain
+radiating to the left arm, onset 3 hours prior to admission. Associated
+diaphoresis and shortness of breath. EKG showed ST depression in V4-V6.
+
+Past Medical History:
+1. Hypertension
+2. Hyperlipidemia
+3. Type 2 Diabetes Mellitus
+
+Social History:
+Non-smoker. Occasional alcohol use. Lives with spouse.
+
+Family History:
+Father with MI at age 60. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.2 HR: 88 BP: 152/94 Resp: 18 O2Sat: 97
+Cardiovascular: Regular rate and rhythm. No murmurs.
+Chest: Clear to auscultation bilaterally.
+
+Pertinent Results:
+Troponin I: 2.4 (elevated)
+BNP: 180
+WBC-9.1 Hgb-11.8 Hct-36.2 Plt-220
+
+Brief Hospital Course:
+Patient admitted for NSTEMI. Heparin drip initiated. Cardiac cath revealed
+70% LAD stenosis; drug-eluting stent placed. Post-procedure course
+uncomplicated. Discharged on dual antiplatelet therapy.
+
+Medications on Admission:
+1. Lisinopril 10 mg daily
+2. Atorvastatin 40 mg nightly
+3. Metformin 1000 mg twice daily
+
+Discharge Medications:
+1. Aspirin 81 mg daily
+2. Clopidogrel 75 mg daily
+3. Lisinopril 10 mg daily
+4. Atorvastatin 80 mg nightly
+5. Metformin 1000 mg twice daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+NSTEMI, coronary artery disease
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Do not stop aspirin or clopidogrel without consulting your cardiologist.
+Call your doctor for chest pain, shortness of breath, or leg swelling.
+
+Followup Instructions:
+Follow up with cardiology in 1 week."
+"d9","10004","20007","DS","1","2150-05-02 12:00:00","2150-05-02 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+sulfa drugs
+
+Attending: ___.
+
+Chief Complaint:
+shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male with known COPD presenting with 3 days of worsening dyspnea
+and productive cough with yellow sputum. Patient reports increased inhaler use
+without improvement. Denies fever. Smoking history: 40 pack-years.
+
+Past Medical History:
+1. COPD, moderate severity
+2. Hypertension
+3. Gastroesophageal reflux disease
+
+Social History:
+Former smoker, quit 5 years ago. No alcohol or illicit drug use. Lives alone.
+
+Family History:
+Mother with asthma. No family history of lung cancer.
+
+Physical Exam:
+Temp: 98.6 HR: 102 BP: 138/82 Resp: 24 O2Sat: 88% on room air
+Chest: Diffuse expiratory wheezes. Prolonged expiratory phase.
+Accessory muscle use noted.
+
+Pertinent Results:
+ABG: pH 7.34 pCO2 52 pO2 58 HCO3 28
+WBC-11.2 (elevated) CXR: hyperinflation, no consolidation
+
+Brief Hospital Course:
+Admitted for COPD exacerbation. Started on systemic steroids, nebulized
+bronchodilators, and supplemental oxygen. Azithromycin added for possible
+infectious trigger. O2 sats improved to 94% on 2L NC. Discharged on
+steroid taper.
+
+Medications on Admission:
+1. Albuterol inhaler PRN
+2. Tiotropium 18 mcg daily
+3. Lisinopril 5 mg daily
+4. Omeprazole 20 mg daily
+
+Discharge Medications:
+1. Prednisone 40 mg daily x 3 days then taper per schedule
+2. Azithromycin 250 mg daily x 3 days (2 doses remaining)
+3. Albuterol inhaler q4h PRN
+4. Tiotropium 18 mcg daily
+5. Lisinopril 5 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+COPD exacerbation
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with mild exertional dyspnea.
+
+Discharge Instructions:
+Complete the full steroid taper. Use albuterol inhaler as needed.
+Return to ER for oxygen saturation below 90%, severe dyspnea, or confusion.
+
+Followup Instructions:
+Follow up with pulmonology in 2 weeks."
+"d10","10005","20008","DS","1","2151-07-22 12:00:00","2151-07-22 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+no known drug allergies
+
+Attending: ___.
+
+Chief Complaint:
+leg swelling and shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old female with history of congestive heart failure (EF 30%) presenting
+with 1 week of progressive bilateral lower extremity edema and orthopnea.
+Weight gain of 8 lbs over past week. Patient reports dietary indiscretion
+with increased sodium intake.
+
+Past Medical History:
+1. Congestive heart failure, EF 30%
+2. Atrial fibrillation
+3. Hypertension
+4. Chronic kidney disease stage 3
+
+Social History:
+Non-smoker. No alcohol. Lives with daughter.
+
+Family History:
+Father with heart failure. Sister with hypertension.
+
+Physical Exam:
+Temp: 98.0 HR: 96 (irregularly irregular) BP: 148/90 Resp: 20 O2Sat: 94%
+JVD present. Bilateral crackles at lung bases.
+2+ pitting edema to knees bilaterally.
+
+Pertinent Results:
+BNP: 1840 (markedly elevated)
+Creatinine 1.8 (baseline 1.5)
+CXR: pulmonary vascular congestion, bilateral pleural effusions
+
+Brief Hospital Course:
+Admitted for decompensated CHF. IV furosemide initiated with good diuretic
+response (negative 3L over 48 hours). Transitioned to oral diuretics.
+Weight at discharge 4 lbs below admission weight. Cardiology consulted.
+
+Medications on Admission:
+1. Furosemide 40 mg daily
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Medications:
+1. Furosemide 80 mg daily (increased)
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+Acute decompensated heart failure
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Weigh yourself daily. If weight increases by more than 2 lbs in a day or
+5 lbs in a week, call your doctor immediately. Restrict sodium to 2g/day.
+
+Followup Instructions:
+Follow up with cardiologist in 3 days."
+"d11","10006","20009","DS","1","2152-09-08 12:00:00","2152-09-08 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+amoxicillin
+
+Attending: ___.
+
+Chief Complaint:
+fever and productive cough
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male presenting with 5 days of productive cough with green sputum,
+fever to 101.8F, and pleuritic chest pain. Patient reports decreased appetite
+and fatigue. No sick contacts. Up to date on vaccinations.
+
+Past Medical History:
+1. Type 2 Diabetes Mellitus
+2. Hypertension
+
+Social History:
+Non-smoker. No alcohol or drug use. Works as a teacher.
+
+Family History:
+No family history of pulmonary disease.
+
+Physical Exam:
+Temp: 101.6 HR: 108 BP: 128/78 Resp: 22 O2Sat: 93% on room air
+Chest: Decreased breath sounds and dullness to percussion at right base.
+Egophony present at right lower lobe.
+
+Pertinent Results:
+WBC-14.8 (elevated) with left shift
+Procalcitonin: 2.1 (elevated)
+CXR: right lower lobe consolidation consistent with pneumonia
+
+Brief Hospital Course:
+Admitted for community-acquired pneumonia. Started on ceftriaxone and
+azithromycin. O2 sats improved on 4L NC. Transitioned to oral antibiotics
+on hospital day 3. Glucose levels monitored and insulin sliding scale used.
+
+Medications on Admission:
+1. Metformin 500 mg twice daily
+2. Lisinopril 10 mg daily
+3. Aspirin 81 mg daily
+
+Discharge Medications:
+1. Levofloxacin 750 mg daily x 5 days (3 doses remaining)
+2. Metformin 500 mg twice daily
+3. Lisinopril 10 mg daily
+4. Aspirin 81 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+Community-acquired pneumonia
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Complete the full antibiotic course. Return for worsening shortness of breath,
+persistent fever, or hemoptysis.
+
+Followup Instructions:
+Follow up with primary care in 1 week for repeat CXR."
+"d12","10006","20010","DS","1","2152-11-18 12:00:00","2152-11-18 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: NEUROLOGY
+
+Allergies:
+codeine
+
+Attending: ___.
+
+Chief Complaint:
+sudden onset left-sided weakness
+
+Major Surgical or Invasive Procedure:
+IV tPA administration
+
+History of Present Illness:
+___ year old female with hypertension and atrial fibrillation brought by EMS
+with sudden onset left arm and leg weakness and facial droop noted 2 hours
+prior to arrival. Last known well time established. NIHSS 12 on arrival.
+CT head showed no hemorrhage. IV tPA administered within window.
+
+Past Medical History:
+1. Atrial fibrillation
+2. Hypertension
+3. Hyperlipidemia
+
+Social History:
+Non-smoker. Rare alcohol use. Retired. Lives alone.
+
+Family History:
+Father with stroke at age 72. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.4 HR: 84 (irregular) BP: 168/96 Resp: 16 O2Sat: 98%
+Neuro: left facial droop, left arm 3/5 strength, left leg 4/5 strength.
+Dysarthria present.
+
+Pertinent Results:
+MRI brain: right MCA territory infarct
+Echo: no thrombus. EF 55%.
+INR: 1.1 (not anticoagulated at time of stroke)
+
+Brief Hospital Course:
+Admitted to stroke unit. tPA given without hemorrhagic complication.
+Deficits improved over 48 hours. Anticoagulation initiated with apixaban.
+Physical therapy and speech therapy consulted. Swallow evaluation passed.
+
+Medications on Admission:
+1. Metoprolol 25 mg twice daily
+2. Amlodipine 5 mg daily
+3. Atorvastatin 20 mg nightly
+
+Discharge Medications:
+1. Apixaban 5 mg twice daily (new)
+2. Aspirin 81 mg daily
+3. Metoprolol 25 mg twice daily
+4. Amlodipine 5 mg daily
+5. Atorvastatin 80 mg nightly (increased)
+
+Discharge Disposition:
+Rehabilitation facility
+
+Discharge Diagnosis:
+Acute ischemic stroke, right MCA territory
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Take apixaban as directed. Do not stop without consulting your neurologist.
+Call 911 immediately for any new weakness, vision changes, or speech difficulty.
+
+Followup Instructions:
+Follow up with neurology in 2 weeks."
+"d13","10007","20011","DS","1","2150-04-11 12:00:00","2150-04-11 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: SURGERY
+
+Allergies:
+morphine
+
+Attending: ___.
+
+Chief Complaint:
+abdominal pain
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+This patient is a ___ year old male who complains of right lower quadrant
+abdominal pain for 2 days. Pain worsens with movement. Patient denies fevers
+or chills. No relief with over-the-counter analgesics.
+
+Past Medical History:
+none
+
+Social History:
+___
+
+Family History:
+NC
+
+Physical Exam:
+Temp: 97.8 HR: 90 BP: 124/86 Resp: 14 O2Sat: 100
+Abdomen: right lower quadrant tenderness without rebound
+
+Pertinent Results:
+WBC-8.9 RBC-5.59 Hgb-12.5 Hct-42.0
+Glucose-99 UreaN-13 Creat-1.0 Na-137 K-4.0 Cl-103 HCO3-22
+
+Brief Hospital Course:
+Patient admitted for observation. Serial abdominal exams performed.
+Pain resolved with conservative management. Diet advanced without issue.
+Discharged home in stable condition.
+
+Medications on Admission:
+none
+
+Discharge Medications:
+none
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+abdominal pain
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Return to ER for worsening pain, fever >101.5F, vomiting, or new symptoms.
+
+Followup Instructions:
+Follow up with primary care physician in 1 week."
+"d14","10008","20012","DS","1","2151-10-25 12:00:00","2151-10-25 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+penicillin
+
+Attending: ___.
+
+Chief Complaint:
+chest pain
+
+Major Surgical or Invasive Procedure:
+cardiac catheterization
+
+History of Present Illness:
+___ year old female with hypertension presenting with substernal chest pain
+radiating to the left arm, onset 3 hours prior to admission. Associated
+diaphoresis and shortness of breath. EKG showed ST depression in V4-V6.
+
+Past Medical History:
+1. Hypertension
+2. Hyperlipidemia
+3. Type 2 Diabetes Mellitus
+
+Social History:
+Non-smoker. Occasional alcohol use. Lives with spouse.
+
+Family History:
+Father with MI at age 60. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.2 HR: 88 BP: 152/94 Resp: 18 O2Sat: 97
+Cardiovascular: Regular rate and rhythm. No murmurs.
+Chest: Clear to auscultation bilaterally.
+
+Pertinent Results:
+Troponin I: 2.4 (elevated)
+BNP: 180
+WBC-9.1 Hgb-11.8 Hct-36.2 Plt-220
+
+Brief Hospital Course:
+Patient admitted for NSTEMI. Heparin drip initiated. Cardiac cath revealed
+70% LAD stenosis; drug-eluting stent placed. Post-procedure course
+uncomplicated. Discharged on dual antiplatelet therapy.
+
+Medications on Admission:
+1. Lisinopril 10 mg daily
+2. Atorvastatin 40 mg nightly
+3. Metformin 1000 mg twice daily
+
+Discharge Medications:
+1. Aspirin 81 mg daily
+2. Clopidogrel 75 mg daily
+3. Lisinopril 10 mg daily
+4. Atorvastatin 80 mg nightly
+5. Metformin 1000 mg twice daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+NSTEMI, coronary artery disease
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Do not stop aspirin or clopidogrel without consulting your cardiologist.
+Call your doctor for chest pain, shortness of breath, or leg swelling.
+
+Followup Instructions:
+Follow up with cardiology in 1 week."
+"d15","10008","20013","DS","1","2151-12-20 12:00:00","2151-12-20 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+sulfa drugs
+
+Attending: ___.
+
+Chief Complaint:
+shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male with known COPD presenting with 3 days of worsening dyspnea
+and productive cough with yellow sputum. Patient reports increased inhaler use
+without improvement. Denies fever. Smoking history: 40 pack-years.
+
+Past Medical History:
+1. COPD, moderate severity
+2. Hypertension
+3. Gastroesophageal reflux disease
+
+Social History:
+Former smoker, quit 5 years ago. No alcohol or illicit drug use. Lives alone.
+
+Family History:
+Mother with asthma. No family history of lung cancer.
+
+Physical Exam:
+Temp: 98.6 HR: 102 BP: 138/82 Resp: 24 O2Sat: 88% on room air
+Chest: Diffuse expiratory wheezes. Prolonged expiratory phase.
+Accessory muscle use noted.
+
+Pertinent Results:
+ABG: pH 7.34 pCO2 52 pO2 58 HCO3 28
+WBC-11.2 (elevated) CXR: hyperinflation, no consolidation
+
+Brief Hospital Course:
+Admitted for COPD exacerbation. Started on systemic steroids, nebulized
+bronchodilators, and supplemental oxygen. Azithromycin added for possible
+infectious trigger. O2 sats improved to 94% on 2L NC. Discharged on
+steroid taper.
+
+Medications on Admission:
+1. Albuterol inhaler PRN
+2. Tiotropium 18 mcg daily
+3. Lisinopril 5 mg daily
+4. Omeprazole 20 mg daily
+
+Discharge Medications:
+1. Prednisone 40 mg daily x 3 days then taper per schedule
+2. Azithromycin 250 mg daily x 3 days (2 doses remaining)
+3. Albuterol inhaler q4h PRN
+4. Tiotropium 18 mcg daily
+5. Lisinopril 5 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+COPD exacerbation
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with mild exertional dyspnea.
+
+Discharge Instructions:
+Complete the full steroid taper. Use albuterol inhaler as needed.
+Return to ER for oxygen saturation below 90%, severe dyspnea, or confusion.
+
+Followup Instructions:
+Follow up with pulmonology in 2 weeks."
+"d16","10009","20014","DS","1","2152-03-23 12:00:00","2152-03-23 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+no known drug allergies
+
+Attending: ___.
+
+Chief Complaint:
+leg swelling and shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old female with history of congestive heart failure (EF 30%) presenting
+with 1 week of progressive bilateral lower extremity edema and orthopnea.
+Weight gain of 8 lbs over past week. Patient reports dietary indiscretion
+with increased sodium intake.
+
+Past Medical History:
+1. Congestive heart failure, EF 30%
+2. Atrial fibrillation
+3. Hypertension
+4. Chronic kidney disease stage 3
+
+Social History:
+Non-smoker. No alcohol. Lives with daughter.
+
+Family History:
+Father with heart failure. Sister with hypertension.
+
+Physical Exam:
+Temp: 98.0 HR: 96 (irregularly irregular) BP: 148/90 Resp: 20 O2Sat: 94%
+JVD present. Bilateral crackles at lung bases.
+2+ pitting edema to knees bilaterally.
+
+Pertinent Results:
+BNP: 1840 (markedly elevated)
+Creatinine 1.8 (baseline 1.5)
+CXR: pulmonary vascular congestion, bilateral pleural effusions
+
+Brief Hospital Course:
+Admitted for decompensated CHF. IV furosemide initiated with good diuretic
+response (negative 3L over 48 hours). Transitioned to oral diuretics.
+Weight at discharge 4 lbs below admission weight. Cardiology consulted.
+
+Medications on Admission:
+1. Furosemide 40 mg daily
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Medications:
+1. Furosemide 80 mg daily (increased)
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+Acute decompensated heart failure
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Weigh yourself daily. If weight increases by more than 2 lbs in a day or
+5 lbs in a week, call your doctor immediately. Restrict sodium to 2g/day.
+
+Followup Instructions:
+Follow up with cardiologist in 3 days."
+"d17","10010","20015","DS","1","2150-08-05 12:00:00","2150-08-05 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+amoxicillin
+
+Attending: ___.
+
+Chief Complaint:
+fever and productive cough
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male presenting with 5 days of productive cough with green sputum,
+fever to 101.8F, and pleuritic chest pain. Patient reports decreased appetite
+and fatigue. No sick contacts. Up to date on vaccinations.
+
+Past Medical History:
+1. Type 2 Diabetes Mellitus
+2. Hypertension
+
+Social History:
+Non-smoker. No alcohol or drug use. Works as a teacher.
+
+Family History:
+No family history of pulmonary disease.
+
+Physical Exam:
+Temp: 101.6 HR: 108 BP: 128/78 Resp: 22 O2Sat: 93% on room air
+Chest: Decreased breath sounds and dullness to percussion at right base.
+Egophony present at right lower lobe.
+
+Pertinent Results:
+WBC-14.8 (elevated) with left shift
+Procalcitonin: 2.1 (elevated)
+CXR: right lower lobe consolidation consistent with pneumonia
+
+Brief Hospital Course:
+Admitted for community-acquired pneumonia. Started on ceftriaxone and
+azithromycin. O2 sats improved on 4L NC. Transitioned to oral antibiotics
+on hospital day 3. Glucose levels monitored and insulin sliding scale used.
+
+Medications on Admission:
+1. Metformin 500 mg twice daily
+2. Lisinopril 10 mg daily
+3. Aspirin 81 mg daily
+
+Discharge Medications:
+1. Levofloxacin 750 mg daily x 5 days (3 doses remaining)
+2. Metformin 500 mg twice daily
+3. Lisinopril 10 mg daily
+4. Aspirin 81 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+Community-acquired pneumonia
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Complete the full antibiotic course. Return for worsening shortness of breath,
+persistent fever, or hemoptysis.
+
+Followup Instructions:
+Follow up with primary care in 1 week for repeat CXR."
+"d18","1","1","DS","1","2150-01-01 12:00:00","2150-01-01 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: NEUROLOGY
+
+Allergies:
+codeine
+
+Attending: ___.
+
+Chief Complaint:
+sudden onset left-sided weakness
+
+Major Surgical or Invasive Procedure:
+IV tPA administration
+
+History of Present Illness:
+___ year old female with hypertension and atrial fibrillation brought by EMS
+with sudden onset left arm and leg weakness and facial droop noted 2 hours
+prior to arrival. Last known well time established. NIHSS 12 on arrival.
+CT head showed no hemorrhage. IV tPA administered within window.
+
+Past Medical History:
+1. Atrial fibrillation
+2. Hypertension
+3. Hyperlipidemia
+
+Social History:
+Non-smoker. Rare alcohol use. Retired. Lives alone.
+
+Family History:
+Father with stroke at age 72. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.4 HR: 84 (irregular) BP: 168/96 Resp: 16 O2Sat: 98%
+Neuro: left facial droop, left arm 3/5 strength, left leg 4/5 strength.
+Dysarthria present.
+
+Pertinent Results:
+MRI brain: right MCA territory infarct
+Echo: no thrombus. EF 55%.
+INR: 1.1 (not anticoagulated at time of stroke)
+
+Brief Hospital Course:
+Admitted to stroke unit. tPA given without hemorrhagic complication.
+Deficits improved over 48 hours. Anticoagulation initiated with apixaban.
+Physical therapy and speech therapy consulted. Swallow evaluation passed.
+
+Medications on Admission:
+1. Metoprolol 25 mg twice daily
+2. Amlodipine 5 mg daily
+3. Atorvastatin 20 mg nightly
+
+Discharge Medications:
+1. Apixaban 5 mg twice daily (new)
+2. Aspirin 81 mg daily
+3. Metoprolol 25 mg twice daily
+4. Amlodipine 5 mg daily
+5. Atorvastatin 80 mg nightly (increased)
+
+Discharge Disposition:
+Rehabilitation facility
+
+Discharge Diagnosis:
+Acute ischemic stroke, right MCA territory
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Take apixaban as directed. Do not stop without consulting your neurologist.
+Call 911 immediately for any new weakness, vision changes, or speech difficulty.
+
+Followup Instructions:
+Follow up with neurology in 2 weeks."
+"d19","1","2","DS","1","2150-01-16 12:00:00","2150-01-16 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: SURGERY
+
+Allergies:
+morphine
+
+Attending: ___.
+
+Chief Complaint:
+abdominal pain
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+This patient is a ___ year old male who complains of right lower quadrant
+abdominal pain for 2 days. Pain worsens with movement. Patient denies fevers
+or chills. No relief with over-the-counter analgesics.
+
+Past Medical History:
+none
+
+Social History:
+___
+
+Family History:
+NC
+
+Physical Exam:
+Temp: 97.8 HR: 90 BP: 124/86 Resp: 14 O2Sat: 100
+Abdomen: right lower quadrant tenderness without rebound
+
+Pertinent Results:
+WBC-8.9 RBC-5.59 Hgb-12.5 Hct-42.0
+Glucose-99 UreaN-13 Creat-1.0 Na-137 K-4.0 Cl-103 HCO3-22
+
+Brief Hospital Course:
+Patient admitted for observation. Serial abdominal exams performed.
+Pain resolved with conservative management. Diet advanced without issue.
+Discharged home in stable condition.
+
+Medications on Admission:
+none
+
+Discharge Medications:
+none
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+abdominal pain
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Return to ER for worsening pain, fever >101.5F, vomiting, or new symptoms.
+
+Followup Instructions:
+Follow up with primary care physician in 1 week."
+"d20","2","3","DS","1","2150-01-01 12:00:00","2150-01-01 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+penicillin
+
+Attending: ___.
+
+Chief Complaint:
+chest pain
+
+Major Surgical or Invasive Procedure:
+cardiac catheterization
+
+History of Present Illness:
+___ year old female with hypertension presenting with substernal chest pain
+radiating to the left arm, onset 3 hours prior to admission. Associated
+diaphoresis and shortness of breath. EKG showed ST depression in V4-V6.
+
+Past Medical History:
+1. Hypertension
+2. Hyperlipidemia
+3. Type 2 Diabetes Mellitus
+
+Social History:
+Non-smoker. Occasional alcohol use. Lives with spouse.
+
+Family History:
+Father with MI at age 60. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.2 HR: 88 BP: 152/94 Resp: 18 O2Sat: 97
+Cardiovascular: Regular rate and rhythm. No murmurs.
+Chest: Clear to auscultation bilaterally.
+
+Pertinent Results:
+Troponin I: 2.4 (elevated)
+BNP: 180
+WBC-9.1 Hgb-11.8 Hct-36.2 Plt-220
+
+Brief Hospital Course:
+Patient admitted for NSTEMI. Heparin drip initiated. Cardiac cath revealed
+70% LAD stenosis; drug-eluting stent placed. Post-procedure course
+uncomplicated. Discharged on dual antiplatelet therapy.
+
+Medications on Admission:
+1. Lisinopril 10 mg daily
+2. Atorvastatin 40 mg nightly
+3. Metformin 1000 mg twice daily
+
+Discharge Medications:
+1. Aspirin 81 mg daily
+2. Clopidogrel 75 mg daily
+3. Lisinopril 10 mg daily
+4. Atorvastatin 80 mg nightly
+5. Metformin 1000 mg twice daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+NSTEMI, coronary artery disease
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Do not stop aspirin or clopidogrel without consulting your cardiologist.
+Call your doctor for chest pain, shortness of breath, or leg swelling.
+
+Followup Instructions:
+Follow up with cardiology in 1 week."
+"d21","2","4","DS","1","2150-01-16 12:00:00","2150-01-16 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+sulfa drugs
+
+Attending: ___.
+
+Chief Complaint:
+shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male with known COPD presenting with 3 days of worsening dyspnea
+and productive cough with yellow sputum. Patient reports increased inhaler use
+without improvement. Denies fever. Smoking history: 40 pack-years.
+
+Past Medical History:
+1. COPD, moderate severity
+2. Hypertension
+3. Gastroesophageal reflux disease
+
+Social History:
+Former smoker, quit 5 years ago. No alcohol or illicit drug use. Lives alone.
+
+Family History:
+Mother with asthma. No family history of lung cancer.
+
+Physical Exam:
+Temp: 98.6 HR: 102 BP: 138/82 Resp: 24 O2Sat: 88% on room air
+Chest: Diffuse expiratory wheezes. Prolonged expiratory phase.
+Accessory muscle use noted.
+
+Pertinent Results:
+ABG: pH 7.34 pCO2 52 pO2 58 HCO3 28
+WBC-11.2 (elevated) CXR: hyperinflation, no consolidation
+
+Brief Hospital Course:
+Admitted for COPD exacerbation. Started on systemic steroids, nebulized
+bronchodilators, and supplemental oxygen. Azithromycin added for possible
+infectious trigger. O2 sats improved to 94% on 2L NC. Discharged on
+steroid taper.
+
+Medications on Admission:
+1. Albuterol inhaler PRN
+2. Tiotropium 18 mcg daily
+3. Lisinopril 5 mg daily
+4. Omeprazole 20 mg daily
+
+Discharge Medications:
+1. Prednisone 40 mg daily x 3 days then taper per schedule
+2. Azithromycin 250 mg daily x 3 days (2 doses remaining)
+3. Albuterol inhaler q4h PRN
+4. Tiotropium 18 mcg daily
+5. Lisinopril 5 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+COPD exacerbation
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with mild exertional dyspnea.
+
+Discharge Instructions:
+Complete the full steroid taper. Use albuterol inhaler as needed.
+Return to ER for oxygen saturation below 90%, severe dyspnea, or confusion.
+
+Followup Instructions:
+Follow up with pulmonology in 2 weeks."
+"d22","2","5","DS","1","2150-01-21 12:00:00","2150-01-21 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: CARDIOLOGY
+
+Allergies:
+no known drug allergies
+
+Attending: ___.
+
+Chief Complaint:
+leg swelling and shortness of breath
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old female with history of congestive heart failure (EF 30%) presenting
+with 1 week of progressive bilateral lower extremity edema and orthopnea.
+Weight gain of 8 lbs over past week. Patient reports dietary indiscretion
+with increased sodium intake.
+
+Past Medical History:
+1. Congestive heart failure, EF 30%
+2. Atrial fibrillation
+3. Hypertension
+4. Chronic kidney disease stage 3
+
+Social History:
+Non-smoker. No alcohol. Lives with daughter.
+
+Family History:
+Father with heart failure. Sister with hypertension.
+
+Physical Exam:
+Temp: 98.0 HR: 96 (irregularly irregular) BP: 148/90 Resp: 20 O2Sat: 94%
+JVD present. Bilateral crackles at lung bases.
+2+ pitting edema to knees bilaterally.
+
+Pertinent Results:
+BNP: 1840 (markedly elevated)
+Creatinine 1.8 (baseline 1.5)
+CXR: pulmonary vascular congestion, bilateral pleural effusions
+
+Brief Hospital Course:
+Admitted for decompensated CHF. IV furosemide initiated with good diuretic
+response (negative 3L over 48 hours). Transitioned to oral diuretics.
+Weight at discharge 4 lbs below admission weight. Cardiology consulted.
+
+Medications on Admission:
+1. Furosemide 40 mg daily
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Medications:
+1. Furosemide 80 mg daily (increased)
+2. Carvedilol 12.5 mg twice daily
+3. Lisinopril 5 mg daily
+4. Warfarin 5 mg daily
+5. Digoxin 0.125 mg daily
+
+Discharge Disposition:
+Home with services
+
+Discharge Diagnosis:
+Acute decompensated heart failure
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Weigh yourself daily. If weight increases by more than 2 lbs in a day or
+5 lbs in a week, call your doctor immediately. Restrict sodium to 2g/day.
+
+Followup Instructions:
+Follow up with cardiologist in 3 days."
+"d23","2","6","DS","1","2150-01-22 12:00:00","2150-01-22 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   M
+
+Service: MEDICINE
+
+Allergies:
+amoxicillin
+
+Attending: ___.
+
+Chief Complaint:
+fever and productive cough
+
+Major Surgical or Invasive Procedure:
+none
+
+History of Present Illness:
+___ year old male presenting with 5 days of productive cough with green sputum,
+fever to 101.8F, and pleuritic chest pain. Patient reports decreased appetite
+and fatigue. No sick contacts. Up to date on vaccinations.
+
+Past Medical History:
+1. Type 2 Diabetes Mellitus
+2. Hypertension
+
+Social History:
+Non-smoker. No alcohol or drug use. Works as a teacher.
+
+Family History:
+No family history of pulmonary disease.
+
+Physical Exam:
+Temp: 101.6 HR: 108 BP: 128/78 Resp: 22 O2Sat: 93% on room air
+Chest: Decreased breath sounds and dullness to percussion at right base.
+Egophony present at right lower lobe.
+
+Pertinent Results:
+WBC-14.8 (elevated) with left shift
+Procalcitonin: 2.1 (elevated)
+CXR: right lower lobe consolidation consistent with pneumonia
+
+Brief Hospital Course:
+Admitted for community-acquired pneumonia. Started on ceftriaxone and
+azithromycin. O2 sats improved on 4L NC. Transitioned to oral antibiotics
+on hospital day 3. Glucose levels monitored and insulin sliding scale used.
+
+Medications on Admission:
+1. Metformin 500 mg twice daily
+2. Lisinopril 10 mg daily
+3. Aspirin 81 mg daily
+
+Discharge Medications:
+1. Levofloxacin 750 mg daily x 5 days (3 doses remaining)
+2. Metformin 500 mg twice daily
+3. Lisinopril 10 mg daily
+4. Aspirin 81 mg daily
+
+Discharge Disposition:
+Home
+
+Discharge Diagnosis:
+Community-acquired pneumonia
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory - Independent.
+
+Discharge Instructions:
+Complete the full antibiotic course. Return for worsening shortness of breath,
+persistent fever, or hemoptysis.
+
+Followup Instructions:
+Follow up with primary care in 1 week for repeat CXR."
+"d24","3","7","DS","1","2150-01-01 12:00:00","2150-01-01 13:00:00","Name:  ___                  Unit No:   ___
+
+Admission Date:  ___              Discharge Date:   ___
+
+Date of Birth:  ___             Sex:   F
+
+Service: NEUROLOGY
+
+Allergies:
+codeine
+
+Attending: ___.
+
+Chief Complaint:
+sudden onset left-sided weakness
+
+Major Surgical or Invasive Procedure:
+IV tPA administration
+
+History of Present Illness:
+___ year old female with hypertension and atrial fibrillation brought by EMS
+with sudden onset left arm and leg weakness and facial droop noted 2 hours
+prior to arrival. Last known well time established. NIHSS 12 on arrival.
+CT head showed no hemorrhage. IV tPA administered within window.
+
+Past Medical History:
+1. Atrial fibrillation
+2. Hypertension
+3. Hyperlipidemia
+
+Social History:
+Non-smoker. Rare alcohol use. Retired. Lives alone.
+
+Family History:
+Father with stroke at age 72. Mother with hypertension.
+
+Physical Exam:
+Temp: 98.4 HR: 84 (irregular) BP: 168/96 Resp: 16 O2Sat: 98%
+Neuro: left facial droop, left arm 3/5 strength, left leg 4/5 strength.
+Dysarthria present.
+
+Pertinent Results:
+MRI brain: right MCA territory infarct
+Echo: no thrombus. EF 55%.
+INR: 1.1 (not anticoagulated at time of stroke)
+
+Brief Hospital Course:
+Admitted to stroke unit. tPA given without hemorrhagic complication.
+Deficits improved over 48 hours. Anticoagulation initiated with apixaban.
+Physical therapy and speech therapy consulted. Swallow evaluation passed.
+
+Medications on Admission:
+1. Metoprolol 25 mg twice daily
+2. Amlodipine 5 mg daily
+3. Atorvastatin 20 mg nightly
+
+Discharge Medications:
+1. Apixaban 5 mg twice daily (new)
+2. Aspirin 81 mg daily
+3. Metoprolol 25 mg twice daily
+4. Amlodipine 5 mg daily
+5. Atorvastatin 80 mg nightly (increased)
+
+Discharge Disposition:
+Rehabilitation facility
+
+Discharge Diagnosis:
+Acute ischemic stroke, right MCA territory
+
+Discharge Condition:
+Mental Status: Clear and coherent.
+Level of Consciousness: Alert and interactive.
+Activity Status: Ambulatory with assistance.
+
+Discharge Instructions:
+Take apixaban as directed. Do not stop without consulting your neurologist.
+Call 911 immediately for any new weakness, vision changes, or speech difficulty.
+
+Followup Instructions:
+Follow up with neurology in 2 weeks."
diff --git a/test-resources/core/mimic4demo/note/radiology.csv b/test-resources/core/mimic4demo/note/radiology.csv
new file mode 100644
index 000000000..c51a6c825
--- /dev/null
+++ b/test-resources/core/mimic4demo/note/radiology.csv
@@ -0,0 +1,265 @@
+note_id,subject_id,hadm_id,note_type,note_seq,charttime,storetime,text
+r1,10001,19999,RR,1,2150-02-15 14:00:00,2150-02-15 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r2,10002,20000,RR,1,2151-01-01 14:00:00,2151-01-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r3,10001,20001,RR,1,2150-03-15 14:00:00,2150-03-15 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r4,10001,20002,RR,1,2150-06-20 14:00:00,2150-06-20 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r5,10002,20003,RR,1,2151-01-10 14:00:00,2151-01-10 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r6,10002,20004,RR,1,2151-04-05 14:00:00,2151-04-05 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r7,10003,20005,RR,1,2152-02-28 14:00:00,2152-02-28 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r8,10003,20006,RR,1,2152-08-10 14:00:00,2152-08-10 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r9,10004,20007,RR,1,2150-05-01 14:00:00,2150-05-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r10,10005,20008,RR,1,2151-07-15 14:00:00,2151-07-15 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r11,10006,20009,RR,1,2152-09-01 14:00:00,2152-09-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r12,10006,20010,RR,1,2152-11-15 14:00:00,2152-11-15 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r13,10007,20011,RR,1,2150-04-10 14:00:00,2150-04-10 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r14,10008,20012,RR,1,2151-10-05 14:00:00,2151-10-05 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r15,10008,20013,RR,1,2151-12-15 14:00:00,2151-12-15 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r16,10009,20014,RR,1,2152-03-20 14:00:00,2152-03-20 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r17,10010,20015,RR,1,2150-08-01 14:00:00,2150-08-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r18,1,1,RR,1,2150-01-01 14:00:00,2150-01-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r19,1,2,RR,1,2150-01-16 14:00:00,2150-01-16 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r20,2,3,RR,1,2150-01-01 14:00:00,2150-01-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r21,2,4,RR,1,2150-01-16 14:00:00,2150-01-16 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r22,2,5,RR,1,2150-01-21 14:00:00,2150-01-21 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r23,2,6,RR,1,2150-01-22 14:00:00,2150-01-22 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
+r24,3,7,RR,1,2150-01-01 14:00:00,2150-01-01 15:00:00,"CLINICAL HISTORY:  Patient fell onto right ribs.  The patient also on anticoagulants.
+
+CHEST AP:
+
+The heart and mediastinum are normal.  The lung fields are clear.  No effusion or pneumothorax is present.
+
+No rib fracture is seen on limited views.
+
+Vascular stent is seen within the lower SVC, position unchanged.
+
+IMPRESSION:  No pneumothorax.  No effusion.  No rib fracture identified."
diff --git a/tests/core/test_bottleneck_transformer.py b/tests/core/test_bottleneck_transformer.py
new file mode 100644
index 000000000..9adbe6432
--- /dev/null
+++ b/tests/core/test_bottleneck_transformer.py
@@ -0,0 +1,108 @@
+import unittest
+from typing import Dict, Type, Union
+
+import torch
+
+from pyhealth.datasets import create_sample_dataset, get_dataloader
+from pyhealth.models import BottleneckTransformer
+from pyhealth.processors.base_processor import FeatureProcessor
+
+
+class TestBottleneckTransformer(unittest.TestCase):
+    """Test cases for the Bottleneck Transformer model."""
+
+    def setUp(self):
+        """Set up test data and model."""
+        self.samples = [
+            {
+                "patient_id": "patient-0",
+                "visit_id": "visit-0",
+                "diagnoses": ["A", "B", "C"],
+                "procedures": ["X", "Y"],
+                "labs": [1.0, 2.0, 3.0],
+                "label": 1,
+            },
+            {
+                "patient_id": "patient-1",
+                "visit_id": "visit-0",
+                "diagnoses": ["D", "E"],
+                "procedures": ["Y"],
+                "labs": [4.0, 5.0, 6.0],
+                "label": 0,
+            },
+        ]
+
+        self.input_schema: Dict[str, Union[str, Type[FeatureProcessor]]] = {
+            "diagnoses": "sequence",
+            "procedures": "sequence",
+            "labs": "tensor",
+        }
+        self.output_schema: Dict[str, Union[str, Type[FeatureProcessor]]] = {
+            "label": "binary"
+        }
+
+        self.dataset = create_sample_dataset(
+            samples=self.samples,
+            input_schema=self.input_schema,
+            output_schema=self.output_schema,
+            dataset_name="test",
+        )
+
+        self.model = BottleneckTransformer(
+            dataset=self.dataset,
+            embedding_dim=128,
+            bottlenecks_n=2,
+            fusion_startidx=1,
+            num_layers=2,
+            heads=2
+        )
+
+    def test_model_initialization(self):
+        """Test that the BottleneckTransformer model initializes correctly."""
+        self.assertIsInstance(self.model, BottleneckTransformer)
+        self.assertEqual(self.model.embedding_dim, 128)
+        self.assertEqual(self.model.bottlenecks_n, 2)
+        self.assertEqual(self.model.num_layers, 2)
+        self.assertEqual(len(self.model.feature_keys), 3)
+        self.assertIn("diagnoses", self.model.feature_keys)
+        self.assertIn("procedures", self.model.feature_keys)
+        self.assertIn("labs", self.model.feature_keys)
+        self.assertEqual(self.model.label_key, "label")
+
+    def test_model_forward(self):
+        """Test that the BottleneckTransformer forward pass works correctly."""
+        train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True)
+        data_batch = next(iter(train_loader))
+
+        with torch.no_grad():
+            ret = self.model(**data_batch)
+
+        self.assertIn("loss", ret)
+        self.assertIn("y_prob", ret)
+        self.assertIn("y_true", ret)
+        self.assertIn("logit", ret)
+
+        self.assertEqual(ret["y_prob"].shape[0], 2)
+        self.assertEqual(ret["y_true"].shape[0], 2)
+        self.assertEqual(ret["logit"].shape[0], 2)
+        self.assertEqual(ret["y_prob"].shape[1], 1)
+        self.assertEqual(ret["y_true"].shape[1], 1)
+        self.assertEqual(ret["logit"].shape[1], 1)
+        self.assertEqual(ret["loss"].dim(), 0)
+
+    def test_model_backward(self):
+        """Test that the BottleneckTransformer backward pass works correctly."""
+        train_loader = get_dataloader(self.dataset, batch_size=2, shuffle=True)
+        data_batch = next(iter(train_loader))
+
+        ret = self.model(**data_batch)
+        ret["loss"].backward()
+
+        has_gradient = any(
+            param.requires_grad and param.grad is not None
+            for param in self.model.parameters()
+        )
+        self.assertTrue(has_gradient, "No parameters have gradients after backward pass")
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_multimodal_mimic4_edge_cases.py b/tests/core/test_multimodal_mimic4_edge_cases.py
new file mode 100644
index 000000000..1f6730307
--- /dev/null
+++ b/tests/core/test_multimodal_mimic4_edge_cases.py
@@ -0,0 +1,75 @@
+from datetime import datetime
+import unittest
+
+
+class _DummyEvent:
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+class _DummyPatient:
+    def __init__(self) -> None:
+        self.patient_id = "p-1"
+        self._admissions = [
+            _DummyEvent(
+                timestamp=datetime(2020, 1, 1, 0, 0, 0),
+                dischtime="malformed-dischtime",
+                hadm_id=101,
+                hospital_expire_flag=0,
+            )
+        ]
+        self._patients = [_DummyEvent(anchor_age=55)]
+
+    def get_events(self, event_type, start=None, end=None, filters=None, return_df=False):
+        if event_type == "patients":
+            return self._patients
+        if event_type == "admissions":
+            return self._admissions
+        if event_type in {"diagnoses_icd", "procedures_icd", "discharge", "radiology"}:
+            return []
+        if event_type == "labevents" and return_df:
+            import polars as pl
+
+            return pl.DataFrame(
+                {
+                    "timestamp": [],
+                    "labevents/itemid": [],
+                    "labevents/storetime": [],
+                    "labevents/valuenum": [],
+                }
+            )
+        return []
+
+
+class TestICDLabsMIMIC4EdgeCases(unittest.TestCase):
+    def test_malformed_dischtime_keeps_temporal_fields_non_empty(self):
+        from pyhealth.processors.stagenet_processor import (
+            StageNetProcessor,
+            StageNetTensorProcessor,
+        )
+        from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4
+
+        task = ICDLabsMIMIC4(window_hours=24)
+        samples = task(_DummyPatient())
+
+        self.assertEqual(len(samples), 1)
+        sample = samples[0]
+
+        self.assertGreater(len(sample["icd_codes"][0]), 0)
+        self.assertGreater(len(sample["labs"][0]), 0)
+        self.assertGreater(len(sample["labs_mask"][0]), 0)
+
+        icd_proc = StageNetProcessor()
+        icd_proc.fit([{"icd_codes": sample["icd_codes"]}], "icd_codes")
+        icd_time, _ = icd_proc.process(sample["icd_codes"])
+        self.assertIsNotNone(icd_time)
+
+        labs_proc = StageNetTensorProcessor()
+        labs_proc.fit([{"labs": sample["labs"]}], "labs")
+        labs_time, _ = labs_proc.process(sample["labs"])
+        self.assertIsNotNone(labs_time)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_notes_labs_mimic4.py b/tests/core/test_notes_labs_mimic4.py
new file mode 100644
index 000000000..4f6c5cb6b
--- /dev/null
+++ b/tests/core/test_notes_labs_mimic4.py
@@ -0,0 +1,409 @@
+"""Unit tests for NotesLabsMIMIC4, admission-section extraction, and ICDLabsMIMIC4 fixes."""
+
+from datetime import datetime
+import unittest
+
+
+class _DummyEvent:
+    def __init__(self, **kwargs):
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+
+
+class _DummyPatientWithNotes:
+    def __init__(self, note_texts=None, icd_codes=None, lab_df=None, vital_df=None) -> None:
+        self.patient_id = "p-1"
+        self._admissions = [
+            _DummyEvent(
+                timestamp=datetime(2020, 1, 1, 0, 0, 0),
+                dischtime="2020-01-03 12:00:00",
+                hadm_id=101,
+                hospital_expire_flag=0,
+            )
+        ]
+        self._patients = [_DummyEvent(anchor_age=55)]
+        self._note_texts = note_texts or []
+        self._icd_codes = icd_codes or []
+        self._lab_df = lab_df
+        self._vital_df = vital_df
+
+    def get_events(self, event_type, start=None, end=None, filters=None, return_df=False):
+        if event_type == "patients":
+            return self._patients
+        if event_type == "admissions":
+            return self._admissions
+        if event_type == "discharge":
+            out = []
+            for text in self._note_texts:
+                out.append(
+                    _DummyEvent(
+                        timestamp=datetime(2020, 1, 3, 12, 0, 0),
+                        text=text,
+                    )
+                )
+            return out
+        if event_type in {"diagnoses_icd", "procedures_icd"}:
+            return self._icd_codes
+        if event_type == "labevents" and return_df:
+            return self._lab_df
+        if event_type == "chartevents" and return_df:
+            return self._vital_df if self._vital_df is not None else pl.DataFrame(
+                {
+                    "timestamp": [],
+                    "chartevents/itemid": [],
+                    "chartevents/storetime": [],
+                    "chartevents/valuenum": [],
+                }
+            )
+        if event_type == "chartevents" and not return_df:
+            return []
+        return []
+
+
+class _DummyPatientMalformedDischtime:
+    def __init__(self) -> None:
+        self.patient_id = "p-2"
+        self._admissions = [
+            _DummyEvent(
+                timestamp=datetime(2020, 1, 1, 0, 0, 0),
+                dischtime="malformed-dischtime",
+                hadm_id=102,
+                hospital_expire_flag=0,
+            )
+        ]
+        self._patients = [_DummyEvent(anchor_age=60)]
+
+    def get_events(self, event_type, start=None, end=None, filters=None, return_df=False):
+        if event_type == "patients":
+            return self._patients
+        if event_type == "admissions":
+            return self._admissions
+        if event_type in {"diagnoses_icd", "procedures_icd", "discharge", "radiology"}:
+            return []
+        if event_type == "labevents" and return_df:
+            import polars as pl
+
+            return pl.DataFrame(
+                {
+                    "timestamp": [],
+                    "labevents/itemid": [],
+                    "labevents/storetime": [],
+                    "labevents/valuenum": [],
+                }
+            )
+        if event_type == "chartevents" and return_df:
+            import polars as pl
+
+            return pl.DataFrame(
+                {
+                    "timestamp": [],
+                    "chartevents/itemid": [],
+                    "chartevents/storetime": [],
+                    "chartevents/valuenum": [],
+                }
+            )
+        return []
+
+
+class TestExtractAdmissionSections(unittest.TestCase):
+    def test_extracts_target_sections(self):
+        from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task
+
+        text = """Chief Complaint:
+Shortness of breath
+
+Past Medical History:
+1. Hypertension
+
+Medications on Admission:
+1. Metoprolol 25 mg PO BID
+
+Discharge Diagnosis:
+Acute MI
+"""
+        result = BaseMultimodalMIMIC4Task._extract_admission_sections(text)
+        self.assertIn("Shortness of breath", result)
+        self.assertIn("Hypertension", result)
+        self.assertIn("Metoprolol", result)
+        self.assertNotIn("Acute MI", result)
+        self.assertIn("[SEP]", result)
+
+    def test_fallback_to_first_1024_when_no_sections(self):
+        from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task
+
+        text = "This is a note with no section headers at all. " * 50
+        result = BaseMultimodalMIMIC4Task._extract_admission_sections(text)
+        self.assertEqual(result, text[:1024])
+
+    def test_case_insensitive_headers(self):
+        from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task
+
+        text = """CHIEF COMPLAINT:
+Chest pain
+
+Past Medical/Surgical History:
+Appendectomy
+"""
+        result = BaseMultimodalMIMIC4Task._extract_admission_sections(text)
+        self.assertIn("Chest pain", result)
+        self.assertIn("Appendectomy", result)
+
+    def test_empty_string_fallback(self):
+        from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task
+
+        result = BaseMultimodalMIMIC4Task._extract_admission_sections("")
+        self.assertEqual(result, "")
+
+
+class TestCollectAdmissionNoteSections(unittest.TestCase):
+    def test_collects_sections_and_returns_time_zero(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4()
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever\n\nPast Medical History:\nDiabetes"]
+        )
+        texts, times = task._collect_admission_note_sections(
+            patient, 101, datetime(2020, 1, 1, 0, 0, 0)
+        )
+        self.assertEqual(len(texts), 1)
+        self.assertIn("Fever", texts[0])
+        self.assertEqual(times, [0.0])
+
+    def test_missing_note_fallback(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4()
+        patient = _DummyPatientWithNotes(note_texts=[])
+        texts, times = task._collect_admission_note_sections(
+            patient, 101, datetime(2020, 1, 1, 0, 0, 0)
+        )
+        self.assertEqual(texts, [""])
+        self.assertEqual(times, [0.0])
+
+    def test_no_time_filter_applied(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4()
+        # Discharge note timestamp is 2020-01-03, well outside any 24h window
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"]
+        )
+        texts, times = task._collect_admission_note_sections(
+            patient, 101, datetime(2020, 1, 1, 0, 0, 0)
+        )
+        self.assertEqual(len(texts), 1)
+        self.assertIn("Fever", texts[0])
+
+
+class TestNotesLabsMIMIC4(unittest.TestCase):
+    def test_default_schema_excludes_icd(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4()
+        self.assertNotIn("icd_codes", task.input_schema)
+        self.assertNotIn("vitals", task.input_schema)
+        self.assertNotIn("vitals_mask", task.input_schema)
+        self.assertIn("admission_note_times", task.input_schema)
+        self.assertIn("labs", task.input_schema)
+        self.assertIn("labs_mask", task.input_schema)
+
+    def test_include_icd_adds_icd_schema(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(include_icd=True)
+        self.assertIn("icd_codes", task.input_schema)
+
+    def test_include_vitals_adds_vitals_schema(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(include_vitals=True)
+        self.assertIn("vitals", task.input_schema)
+        self.assertIn("vitals_mask", task.input_schema)
+        self.assertNotIn("icd_codes", task.input_schema)
+
+    def test_include_vitals_and_icd_adds_both(self):
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(include_icd=True, include_vitals=True)
+        self.assertIn("vitals", task.input_schema)
+        self.assertIn("vitals_mask", task.input_schema)
+        self.assertIn("icd_codes", task.input_schema)
+
+    def test_output_structure_with_vitals(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(window_hours=24, include_vitals=True)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1, 2, 0, 0)],
+                "labevents/itemid": ["50824"],
+                "labevents/storetime": ["2020-01-01 02:00:00"],
+                "labevents/valuenum": [138.0],
+            }
+        )
+        vital_df = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1, 1, 0, 0)],
+                "chartevents/itemid": ["220045"],  # HeartRate
+                "chartevents/storetime": ["2020-01-01 01:00:00"],
+                "chartevents/valuenum": [80.0],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"],
+            lab_df=lab_df,
+            vital_df=vital_df,
+        )
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        sample = samples[0]
+        self.assertIn("vitals", sample)
+        self.assertIn("vitals_mask", sample)
+        self.assertNotIn("icd_codes", sample)
+        vital_times, vital_values = sample["vitals"]
+        self.assertGreater(len(vital_times), 0)
+
+    def test_vitals_fallback_when_empty(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(window_hours=24, include_vitals=True)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [],
+                "labevents/itemid": [],
+                "labevents/storetime": [],
+                "labevents/valuenum": [],
+            }
+        )
+        vital_df = pl.DataFrame(
+            {
+                "timestamp": [],
+                "chartevents/itemid": [],
+                "chartevents/storetime": [],
+                "chartevents/valuenum": [],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"],
+            lab_df=lab_df,
+            vital_df=vital_df,
+        )
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        vital_times, vital_values = samples[0]["vitals"]
+        self.assertEqual(len(vital_times), 1)
+        self.assertEqual(len(vital_values[0]), len(task.VITAL_CATEGORY_NAMES))
+
+    def test_output_structure_no_icd(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(window_hours=24)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [datetime(2020, 1, 1, 2, 0, 0)],
+                "labevents/itemid": ["50824"],  # Sodium
+                "labevents/storetime": ["2020-01-01 02:00:00"],
+                "labevents/valuenum": [138.0],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"],
+            lab_df=lab_df,
+        )
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        sample = samples[0]
+        self.assertIn("admission_note_times", sample)
+        self.assertIn("labs", sample)
+        self.assertIn("labs_mask", sample)
+        self.assertNotIn("icd_codes", sample)
+        self.assertEqual(sample["mortality"], 0)
+
+    def test_output_structure_with_icd(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(window_hours=24, include_icd=True)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [],
+                "labevents/itemid": [],
+                "labevents/storetime": [],
+                "labevents/valuenum": [],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"],
+            icd_codes=[_DummyEvent(icd_code="I21")],
+            lab_df=lab_df,
+        )
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        sample = samples[0]
+        self.assertIn("icd_codes", sample)
+
+    def test_mortality_label_positive(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import NotesLabsMIMIC4
+
+        task = NotesLabsMIMIC4(window_hours=24)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [],
+                "labevents/itemid": [],
+                "labevents/storetime": [],
+                "labevents/valuenum": [],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=["Chief Complaint:\nFever"],
+            lab_df=lab_df,
+        )
+        patient._admissions[0].hospital_expire_flag = 1
+        samples = task(patient)
+        self.assertEqual(samples[0]["mortality"], 1)
+
+
+class TestICDLabsMIMIC4Fixes(unittest.TestCase):
+    def test_malformed_dischtime_does_not_drop_admission(self):
+        from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4
+
+        task = ICDLabsMIMIC4(window_hours=24)
+        patient = _DummyPatientMalformedDischtime()
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        sample = samples[0]
+        self.assertGreater(len(sample["icd_codes"][0]), 0)
+        self.assertGreater(len(sample["labs"][0]), 0)
+        self.assertGreater(len(sample["labs_mask"][0]), 0)
+
+    def test_missing_icd_code_uses_text_token(self):
+        import polars as pl
+        from pyhealth.tasks.multimodal_mimic4 import ICDLabsMIMIC4
+
+        task = ICDLabsMIMIC4(window_hours=24)
+        lab_df = pl.DataFrame(
+            {
+                "timestamp": [],
+                "labevents/itemid": [],
+                "labevents/storetime": [],
+                "labevents/valuenum": [],
+            }
+        )
+        patient = _DummyPatientWithNotes(
+            note_texts=[],
+            icd_codes=[],  # no ICD codes
+            lab_df=lab_df,
+        )
+        samples = task(patient)
+        self.assertEqual(len(samples), 1)
+        _, icd_visits = samples[0]["icd_codes"]
+        self.assertEqual(icd_visits, [[""]])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_stagenet_time_tensor.py b/tests/core/test_stagenet_time_tensor.py
new file mode 100644
index 000000000..2813ed73b
--- /dev/null
+++ b/tests/core/test_stagenet_time_tensor.py
@@ -0,0 +1,27 @@
+import unittest
+
+
+class TestStageNetTimeTensor(unittest.TestCase):
+    def test_stagenet_processor_time_never_none(self):
+        from pyhealth.processors.stagenet_processor import StageNetProcessor
+
+        proc = StageNetProcessor()
+        proc.fit([{"codes": (None, ["A", "B", "C"])}], "codes")
+        time_tensor, value_tensor = proc.process((None, ["A", "B", "C"]))
+
+        self.assertIsNotNone(time_tensor)
+        self.assertEqual(time_tensor.shape[0], value_tensor.shape[0])
+
+    def test_stagenet_tensor_processor_time_never_none(self):
+        from pyhealth.processors.stagenet_processor import StageNetTensorProcessor
+
+        proc = StageNetTensorProcessor()
+        proc.fit([{"labs": (None, [[1.0, 2.0], [3.0, 4.0]])}], "labs")
+        time_tensor, value_tensor = proc.process((None, [[1.0, 2.0], [3.0, 4.0]]))
+
+        self.assertIsNotNone(time_tensor)
+        self.assertEqual(time_tensor.shape[0], value_tensor.shape[0])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_time_image_processor.py b/tests/core/test_time_image_processor.py
index 51467aa5c..b1eb46cc2 100644
--- a/tests/core/test_time_image_processor.py
+++ b/tests/core/test_time_image_processor.py
@@ -281,13 +281,14 @@ def test_size_set_after_process(self):
     def test_repr(self):
         """__repr__ contains key parameter values."""
         proc = TimeImageProcessor(
-            image_size=128, max_images=5, mode="L"
+            image_size=128, max_images=5, mode="L", padding=""
         )
         r = repr(proc)
         self.assertIn("TimeImageProcessor", r)
         self.assertIn("image_size=128", r)
         self.assertIn("max_images=5", r)
         self.assertIn("mode=L", r)
+        self.assertIn("padding=", r)
 
     # ---- Tensor properties ----
 
@@ -310,6 +311,18 @@ def test_timestamp_dtype_is_float32(self):
 
         self.assertEqual(timestamps.dtype, torch.float32)
 
+    # ---- padding ----
+
+    def test_padding_returns_zero_tensor(self):
+        """Padding path returns a zero tensor instead of loading."""
+        proc = TimeImageProcessor(image_size=32, padding="")
+        images, timestamps, tag = proc.process(
+            (["", self.rgb_paths[0]], [0.0, 1.0])
+        )
+
+        self.assertEqual(images.shape, (2, 3, 32, 32))
+        self.assertTrue(torch.all(images[0] == 0))
+
 
 if __name__ == "__main__":
     unittest.main()
\ No newline at end of file
diff --git a/tests/core/test_unified_e2e_mimic4.py b/tests/core/test_unified_e2e_mimic4.py
new file mode 100644
index 000000000..884544573
--- /dev/null
+++ b/tests/core/test_unified_e2e_mimic4.py
@@ -0,0 +1,154 @@
+from pathlib import Path
+import tempfile
+import unittest
+from datetime import datetime
+
+import numpy as np
+
+from pyhealth.datasets import MIMIC4Dataset, get_dataloader
+from pyhealth.models import MLP, RNN, UnifiedMultimodalEmbeddingModel
+from pyhealth.tasks import ICDLabsMIMIC4
+from pyhealth.trainer import Trainer
+
+
+class TestUnifiedE2EMIMIC4(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.cache_dir = tempfile.TemporaryDirectory()
+        cls.sample_dataset = None
+        cls.small_dataset = None
+        cls.ehr_root = str(
+            Path(__file__).parent.parent.parent
+            / "test-resources"
+            / "core"
+            / "mimic4demo"
+        )
+
+        try:
+            base_dataset = MIMIC4Dataset(
+                ehr_root=cls.ehr_root,
+                ehr_tables=["diagnoses_icd", "procedures_icd", "labevents"],
+                cache_dir=cls.cache_dir.name,
+                num_workers=1,
+            )
+            task = ICDLabsMIMIC4()
+            cls.sample_dataset = base_dataset.set_task(task, num_workers=1)
+            if len(cls.sample_dataset) == 0:
+                raise unittest.SkipTest(
+                    "ICDLabsMIMIC4 produced no demo samples."
+                )
+            max_samples = min(16, len(cls.sample_dataset))
+            cls.small_dataset = cls.sample_dataset.subset(list(range(max_samples)))
+        except Exception as exc:
+            raise unittest.SkipTest(
+                f"Skipping MIMIC4 unified E2E integration test due to dataset backend issue: {exc}"
+            ) from exc
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.sample_dataset is not None:
+            cls.sample_dataset.close()
+        if cls.small_dataset is not None:
+            cls.small_dataset.close()
+        cls.cache_dir.cleanup()
+
+    def test_task_outputs_expected_fields(self):
+        sample = self.sample_dataset[0]
+        for key in [
+            "patient_id",
+            "icd_codes",
+            "labs",
+            "labs_mask",
+            "mortality",
+            "window_start",
+            "window_end",
+        ]:
+            self.assertIn(key, sample)
+        self.assertIsInstance(sample["icd_codes"], tuple)
+        self.assertIsInstance(sample["labs"], tuple)
+
+    def test_deterministic_window_boundaries(self):
+        from pyhealth.tasks.multimodal_mimic4 import BaseMultimodalMIMIC4Task
+
+        class _DummyTask(BaseMultimodalMIMIC4Task):
+            pass
+
+        class _Admission:
+            def __init__(self, timestamp, dischtime):
+                self.timestamp = timestamp
+                self.dischtime = dischtime
+
+        global_start = datetime(2020, 1, 1, 0, 0, 0)
+        admissions = [
+            _Admission(global_start, "2020-01-01 12:00:00"),
+            _Admission(datetime(2020, 1, 2, 0, 0, 0), "2020-01-02 12:00:00"),
+        ]
+
+        task_windowed = _DummyTask(window_hours=12)
+        pairs = {task_windowed._compute_effective_window(admissions) for _ in range(20)}
+        self.assertEqual(len(pairs), 1)
+        only_pair = next(iter(pairs))
+        self.assertEqual(only_pair[0], global_start)
+        self.assertEqual(only_pair[1], datetime(2020, 1, 1, 12, 0, 0))
+
+        task_full = _DummyTask(window_hours=None)
+        full_pairs = {
+            task_full._compute_effective_window(admissions) for _ in range(20)
+        }
+        self.assertEqual(len(full_pairs), 1)
+
+    def _run_and_check_model(self, model_type: str):
+        unified = UnifiedMultimodalEmbeddingModel(
+            processors=self.small_dataset.input_processors,
+            embedding_dim=32,
+        )
+        if model_type == "mlp":
+            model = MLP(
+                dataset=self.small_dataset,
+                embedding_dim=32,
+                hidden_dim=32,
+                unified_embedding=unified,
+            )
+        else:
+            model = RNN(
+                dataset=self.small_dataset,
+                embedding_dim=32,
+                hidden_dim=32,
+                unified_embedding=unified,
+                rnn_type="GRU",
+                num_layers=1,
+                dropout=0.0,
+            )
+
+        loader = get_dataloader(self.small_dataset, batch_size=4, shuffle=False)
+        batch = next(iter(loader))
+        forward_out = model(**batch)
+
+        self.assertIn("y_prob", forward_out)
+        self.assertIn("loss", forward_out)
+        self.assertGreater(forward_out["y_prob"].shape[0], 0)
+
+        trainer = Trainer(
+            model=model,
+            metrics=["accuracy"],
+            device="cpu",
+            enable_logging=False,
+        )
+        y_true, y_prob, _, patient_ids = trainer.inference(
+            loader, return_patient_ids=True
+        )
+
+        self.assertEqual(y_true.shape[0], y_prob.shape[0])
+        self.assertEqual(y_prob.shape[0], len(patient_ids))
+        self.assertTrue(np.all(y_prob >= 0.0))
+        self.assertTrue(np.all(y_prob <= 1.0))
+
+    def test_unified_mlp_e2e_prediction(self):
+        self._run_and_check_model(model_type="mlp")
+
+    def test_unified_rnn_e2e_prediction(self):
+        self._run_and_check_model(model_type="rnn")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tests/core/test_vision_embedding.py b/tests/core/test_vision_embedding.py
index 39d3e0c66..7823cc34d 100644
--- a/tests/core/test_vision_embedding.py
+++ b/tests/core/test_vision_embedding.py
@@ -19,7 +19,7 @@
 
 from pyhealth.datasets import create_sample_dataset
 from pyhealth.datasets.utils import get_dataloader
-from pyhealth.models.vision_embedding import VisionEmbeddingModel
+from pyhealth.models.embedding import VisionEmbeddingModel
 
 
 class TestVisionEmbeddingModel(unittest.TestCase):
diff --git a/tests/datasets/test_collate.py b/tests/datasets/test_collate.py
new file mode 100644
index 000000000..2f2a15c42
--- /dev/null
+++ b/tests/datasets/test_collate.py
@@ -0,0 +1,133 @@
+"""Tests for the collate_temporal function."""
+import torch
+from torch.nn.utils.rnn import pad_sequence
+from pyhealth.datasets.collate import collate_temporal
+
+
+def test_collate_temporal_tuple_handling():
+    """Test that collate_temporal correctly handles tuple outputs from processors like TupleTimeTextProcessor."""
+    # Simulate a batch of two samples where a field returns a tuple (text_list, time_tensor, tag)
+    # This mimics the output of TupleTimeTextProcessor without tokenizer.
+    batch = [
+        {
+            "patient_id": ["p1", "p2"],
+            "notes": (["Note 1", "Note 2"], torch.tensor([0.0, 10.0]), "note"),
+            "label": 0,
+        },
+        {
+            "patient_id": ["p3"],
+            "notes": (["Note 3"], torch.tensor([5.0]), "note"),
+            "label": 1,
+        },
+    ]
+
+    collated = collate_temporal(batch)
+
+    # Check that the notes field is a tuple of three elements: padded text list, padded time tensor, and list of tags
+    assert isinstance(collated["notes"], tuple)
+    assert len(collated["notes"]) == 3
+
+    # First element: list of text lists (should not be converted to tensor)
+    assert collated["notes"][0] == [["Note 1", "Note 2"], ["Note 3"]]
+
+    # Second element: time tensor padded to longest sequence
+    expected_times = pad_sequence(
+        [torch.tensor([0.0, 10.0]), torch.tensor([5.0])], batch_first=True, padding_value=0.0
+    )
+    assert torch.equal(collated["notes"][1], expected_times)
+
+    # Third element: list of tags
+    assert collated["notes"][2] == ["note", "note"]
+
+
+def test_collate_temporal_tuple_with_tokenizer():
+    """Test collation when tuple contains tensors (tokenized text output)."""
+    # Simulate tokenized output: (input_ids, attention_mask, token_type_ids, time_tensor, tag)
+    batch = [
+        {
+            "patient_id": ["p1"],
+            "notes": (
+                torch.tensor([[101, 2023, 2003, 102], [101, 2003, 2023, 102]]),  # input_ids (2,4)
+                torch.tensor([[1, 1, 1, 1], [1, 1, 1, 1]]),                   # attention_mask (2,4)
+                torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0]]),                   # token_type_ids (2,4)
+                torch.tensor([0.0, 10.0]),                                    # time (2,)
+                "note",
+            ),
+            "label": 0,
+        },
+        {
+            "patient_id": ["p2"],
+            "notes": (
+                torch.tensor([[101, 2023, 102, 0], [0, 0, 0, 0]]),            # input_ids (2,4) - padded
+                torch.tensor([[1, 1, 1, 0], [0, 0, 0, 0]]),                 # attention_mask (2,4)
+                torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0]]),                 # token_type_ids (2,4)
+                torch.tensor([5.0, 0.0]),                                     # time (2,)
+                "note",
+            ),
+            "label": 1,
+        },
+    ]
+
+    collated = collate_temporal(batch)
+
+    assert isinstance(collated["notes"], tuple)
+    assert len(collated["notes"]) == 5
+
+    # input_ids should be stacked (both sequences now length 4)
+    expected_input_ids = torch.stack([
+        torch.tensor([[101, 2023, 2003, 102], [101, 2003, 2023, 102]]),  # (2,4)
+        torch.tensor([[101, 2023, 102, 0], [0, 0, 0, 0]])                 # (2,4) padded
+    ])
+    assert torch.equal(collated["notes"][0], expected_input_ids)
+
+    # attention_mask similarly stacked
+    expected_attention_mask = torch.stack([
+        torch.tensor([[1, 1, 1, 1], [1, 1, 1, 1]]),  # (2,4)
+        torch.tensor([[1, 1, 1, 0], [0, 0, 0, 0]])    # (2,4) padded
+    ])
+    assert torch.equal(collated["notes"][1], expected_attention_mask)
+
+    # token_type_ids stacked
+    expected_token_type_ids = torch.stack([
+        torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0]]),  # (2,4)
+        torch.tensor([[0, 0, 0, 0], [0, 0, 0, 0]])    # (2,4) padded
+    ])
+    assert torch.equal(collated["notes"][2], expected_token_type_ids)
+
+    # time tensor stacked
+    expected_time = torch.stack([
+        torch.tensor([0.0, 10.0]),    # (2,)
+        torch.tensor([5.0, 0.0])      # (2,) padded
+    ])
+    assert torch.equal(collated["notes"][3], expected_time)
+
+    # tags
+    assert collated["notes"][4] == ["note", "note"]
+
+
+def test_collate_temporal_mixed():
+    """Test collation with a mix of tensor, tuple, and scalar fields."""
+    batch = [
+        {
+            "id": 1,
+            "value": torch.tensor([1.0, 2.0]),
+            "meta": (["a", "b"], torch.tensor([0.0, 1.0])),
+        },
+        {
+            "id": 2,
+            "value": torch.tensor([3.0]),
+            "meta": (["c"], torch.tensor([2.0])),
+        },
+    ]
+
+    collated = collate_temporal(batch)
+
+    assert torch.equal(collated["id"], torch.tensor([1, 2]))
+    # value tensor padded to length 2
+    assert torch.equal(collated["value"], torch.tensor([[1.0, 2.0], [3.0, 0.0]]))
+    # meta tuple: first element list of lists, second element padded time tensor
+    assert collated["meta"][0] == [["a", "b"], ["c"]]
+    assert torch.equal(
+        collated["meta"][1],
+        pad_sequence([torch.tensor([0.0, 1.0]), torch.tensor([2.0])], batch_first=True, padding_value=0.0),
+    )
\ No newline at end of file
diff --git a/tests/test_text_embedding.py b/tests/test_text_embedding.py
index 10e97f8f5..343024ed6 100644
--- a/tests/test_text_embedding.py
+++ b/tests/test_text_embedding.py
@@ -1,7 +1,7 @@
 import torch
 import pytest
 import warnings
-from pyhealth.models.text_embedding import TextEmbedding
+from pyhealth.models.embedding import TextEmbedding
 
 
 def test_text_embedding_initialization():
diff --git a/tests/test_tuple_time_text_processor.py b/tests/test_tuple_time_text_processor.py
index 7a8dfd5b5..b5fbd48e3 100644
--- a/tests/test_tuple_time_text_processor.py
+++ b/tests/test_tuple_time_text_processor.py
@@ -29,3 +29,18 @@ def test_tuple_time_text_processor():
     from pyhealth.processors import get_processor
     ProcessorClass = get_processor("tuple_time_text")
     assert ProcessorClass is TupleTimeTextProcessor
+
+
+def test_tuple_time_text_processor_empty_input_fallback():
+    """Empty or whitespace-only text lists should not crash processing."""
+    processor = TupleTimeTextProcessor(type_tag="clinical_note")
+
+    texts = ["   ", None, ""]
+    time_diffs = [1.0, 2.0, 3.0]
+    result_texts, time_tensor, tag = processor.process((texts, time_diffs))
+
+    assert result_texts == ["[MISSING_TEXT]"]
+    assert isinstance(time_tensor, torch.Tensor)
+    assert time_tensor.shape == (1,)
+    assert torch.equal(time_tensor, torch.tensor([0.0]))
+    assert tag == "clinical_note"
diff --git a/tests/test_tuple_time_text_tokenizer.py b/tests/test_tuple_time_text_tokenizer.py
index 3a5eeb599..1024e2a70 100644
--- a/tests/test_tuple_time_text_tokenizer.py
+++ b/tests/test_tuple_time_text_tokenizer.py
@@ -121,6 +121,27 @@ def test_tokenizer_integration_in_pyhealth_workflow():
     assert len(result) == 5        # (input_ids, mask, type_ids, time, tag)
 
 
+@pytest.mark.skipif(not TRANSFORMERS_AVAILABLE, reason="Transformers not installed")
+def test_tuple_time_text_processor_with_tokenizer_empty_input_fallback():
+    """Tokenizer mode should handle empty/malformed note batches gracefully."""
+    processor = TupleTimeTextProcessor(
+        tokenizer_model="prajjwal1/bert-tiny",
+        max_length=8,
+    )
+
+    # Empty texts and invalid timestamps should fallback to one missing token.
+    input_ids, attention_mask, token_type_ids, time_tensor, tag = processor.process(
+        (["", "   ", None], ["bad", None, "nan"])
+    )
+
+    assert input_ids.shape == (1, 8)
+    assert attention_mask.shape == (1, 8)
+    assert token_type_ids.shape == (1, 8)
+    assert time_tensor.shape == (1,)
+    assert torch.equal(time_tensor, torch.tensor([0.0]))
+    assert tag == "note"
+
+
 @pytest.mark.skipif(not TRANSFORMERS_AVAILABLE, reason="Transformers not installed")
 def test_tuple_in_schema_canonical_form():
     """Canonical usage: declare processor via ('tuple_time_text', kwargs) in input_schema.
@@ -195,4 +216,3 @@ def test_tuple_in_schema_canonical_form():
         out = model(**batch)
     assert "loss" in out
     assert out["y_prob"].shape == (2, 1)
-
diff --git a/tests/test_unified_multimodal.py b/tests/test_unified_multimodal.py
index 01bc194e6..9cac102e7 100644
--- a/tests/test_unified_multimodal.py
+++ b/tests/test_unified_multimodal.py
@@ -2,49 +2,73 @@
 collate_temporal helper, and UnifiedMultimodalEmbeddingModel.
 
 Run with:
-    TOKENIZERS_PARALLELISM=false pytest tests/test_unified_multimodal.py -v
+    TOKENIZERS_PARALLELISM=false python tests/test_unified_multimodal.py
 """
+
 import math
+import unittest
 from datetime import datetime, timedelta
+from unittest.mock import patch
 
-import pytest
 import torch
 import numpy as np
 
 
 # ── 1. TemporalFeatureProcessor ABC & ModalityType ────────────────────────────
 
+
 def test_modality_type_values():
     from pyhealth.processors import ModalityType
-    assert ModalityType.CODE    == "code"
-    assert ModalityType.TEXT    == "text"
-    assert ModalityType.IMAGE   == "image"
+
+    assert ModalityType.CODE == "code"
+    assert ModalityType.TEXT == "text"
+    assert ModalityType.IMAGE == "image"
     assert ModalityType.NUMERIC == "numeric"
 
 
 def test_stagenet_is_temporal():
-    from pyhealth.processors import StageNetProcessor, TemporalFeatureProcessor, ModalityType
+    from pyhealth.processors import (
+        StageNetProcessor,
+        TemporalFeatureProcessor,
+        ModalityType,
+    )
+
     p = StageNetProcessor()
     assert isinstance(p, TemporalFeatureProcessor)
     assert p.modality() == ModalityType.CODE
 
 
 def test_stagenet_tensor_is_temporal():
-    from pyhealth.processors import StageNetTensorProcessor, TemporalFeatureProcessor, ModalityType
+    from pyhealth.processors import (
+        StageNetTensorProcessor,
+        TemporalFeatureProcessor,
+        ModalityType,
+    )
+
     p = StageNetTensorProcessor()
     assert isinstance(p, TemporalFeatureProcessor)
     assert p.modality() == ModalityType.NUMERIC
 
 
 def test_tuple_time_text_is_temporal():
-    from pyhealth.processors import TupleTimeTextProcessor, TemporalFeatureProcessor, ModalityType
+    from pyhealth.processors import (
+        TupleTimeTextProcessor,
+        TemporalFeatureProcessor,
+        ModalityType,
+    )
+
     p = TupleTimeTextProcessor()
     assert isinstance(p, TemporalFeatureProcessor)
     assert p.modality() == ModalityType.TEXT
 
 
 def test_time_image_is_temporal():
-    from pyhealth.processors import TimeImageProcessor, TemporalFeatureProcessor, ModalityType
+    from pyhealth.processors import (
+        TimeImageProcessor,
+        TemporalFeatureProcessor,
+        ModalityType,
+    )
+
     p = TimeImageProcessor()
     assert isinstance(p, TemporalFeatureProcessor)
     assert p.modality() == ModalityType.IMAGE
@@ -52,8 +76,10 @@ def test_time_image_is_temporal():
 
 # ── 2. StageNetProcessor.process_temporal() ───────────────────────────────────
 
+
 def test_stagenet_process_temporal():
     from pyhealth.processors import StageNetProcessor
+
     samples = [{"codes": (None, ["A", "B", "C"])}]
     p = StageNetProcessor()
     p.fit(samples, "codes")
@@ -63,13 +89,14 @@ def test_stagenet_process_temporal():
 
     assert set(out.keys()) == {"value", "time"}
     assert out["value"].dtype == torch.long
-    assert out["time"].dtype  == torch.float32
+    assert out["time"].dtype == torch.float32
     assert out["value"].shape == (3,)
-    assert out["time"].shape  == (3,)
+    assert out["time"].shape == (3,)
 
 
 def test_stagenet_tensor_process_temporal():
     from pyhealth.processors import StageNetTensorProcessor
+
     samples = [{"vitals": ([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]])}]
     p = StageNetTensorProcessor()
     p.fit(samples, "vitals")
@@ -77,35 +104,37 @@ def test_stagenet_tensor_process_temporal():
     out = p.process_temporal(([0.0, 1.0], [[1.0, 2.0], [3.0, 4.0]]))
     assert set(out.keys()) == {"value", "time"}
     assert out["value"].shape == (2, 2)
-    assert out["time"].shape  == (2,)
+    assert out["time"].shape == (2,)
     assert p.value_dim() == 2
     assert p.modality().value == "numeric"
 
 
 # ── 3. TemporalTimeseriesProcessor ────────────────────────────────────────────
 
+
 def test_temporal_timeseries_basic():
     from pyhealth.processors import TemporalTimeseriesProcessor
 
     p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=2))
     ts = [
-        datetime(2023, 1, 1,  0),
-        datetime(2023, 1, 1,  4),
-        datetime(2023, 1, 1,  8),
+        datetime(2023, 1, 1, 0),
+        datetime(2023, 1, 1, 4),
+        datetime(2023, 1, 1, 8),
     ]
     val = np.array([[120.0, 80.0], [115.0, 78.0], [118.0, 82.0]])
 
     out = p.process((ts, val))
     # 8 h window / 2 h step + 1 = 5 steps
     assert out["value"].shape == (5, 2)
-    assert out["time"].shape  == (5,)
+    assert out["time"].shape == (5,)
     # Times should be [0, 2, 4, 6, 8]
-    expected_times = torch.tensor([0., 2., 4., 6., 8.])
+    expected_times = torch.tensor([0.0, 2.0, 4.0, 6.0, 8.0])
     assert torch.allclose(out["time"], expected_times)
 
 
 def test_temporal_timeseries_fit():
     from pyhealth.processors import TemporalTimeseriesProcessor
+
     p = TemporalTimeseriesProcessor()
     ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 1)]
     val = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
@@ -117,29 +146,35 @@ def test_temporal_timeseries_fit():
 
 def test_temporal_timeseries_imputation():
     from pyhealth.processors import TemporalTimeseriesProcessor
+
     p = TemporalTimeseriesProcessor(sampling_rate=timedelta(hours=1))
     ts = [datetime(2023, 1, 1, 0), datetime(2023, 1, 1, 2)]  # gap at h=1
     val = np.array([[10.0], [20.0]])
     out = p.process((ts, val))
     # 3 steps: h=0 → 10, h=1 → forward-filled to 10, h=2 → 20
     assert out["value"].shape == (3, 1)
-    assert float(out["value"][1, 0]) == pytest.approx(10.0)
+    assert math.isclose(float(out["value"][1, 0]), 10.0, rel_tol=1e-6)
 
 
 # ── 4. collate_temporal ───────────────────────────────────────────────────────
 
+
 def test_collate_temporal_basic():
     from pyhealth.datasets.collate import collate_temporal
 
     batch = [
         {
-            "codes": {"value": torch.tensor([1, 2, 3], dtype=torch.long),
-                      "time":  torch.tensor([0., 1., 2.])},
+            "codes": {
+                "value": torch.tensor([1, 2, 3], dtype=torch.long),
+                "time": torch.tensor([0.0, 1.0, 2.0]),
+            },
             "label": torch.tensor(1),
         },
         {
-            "codes": {"value": torch.tensor([4, 5, 3], dtype=torch.long),
-                      "time":  torch.tensor([0.5, 1.5, 2.5])},
+            "codes": {
+                "value": torch.tensor([4, 5, 3], dtype=torch.long),
+                "time": torch.tensor([0.5, 1.5, 2.5]),
+            },
             "label": torch.tensor(0),
         },
     ]
@@ -147,7 +182,7 @@ def test_collate_temporal_basic():
     collated = collate_temporal(batch)
 
     assert collated["codes"]["value"].shape == (2, 3)
-    assert collated["codes"]["time"].shape  == (2, 3)
+    assert collated["codes"]["time"].shape == (2, 3)
     assert collated["label"].shape == (2,)
 
 
@@ -156,10 +191,18 @@ def test_collate_temporal_variable_length():
     from pyhealth.datasets.collate import collate_temporal
 
     batch = [
-        {"codes": {"value": torch.tensor([1, 2], dtype=torch.long),
-                   "time":  torch.tensor([0., 1.])}},
-        {"codes": {"value": torch.tensor([3, 4, 5], dtype=torch.long),
-                   "time":  torch.tensor([0., 1., 2.])}},
+        {
+            "codes": {
+                "value": torch.tensor([1, 2], dtype=torch.long),
+                "time": torch.tensor([0.0, 1.0]),
+            }
+        },
+        {
+            "codes": {
+                "value": torch.tensor([3, 4, 5], dtype=torch.long),
+                "time": torch.tensor([0.0, 1.0, 2.0]),
+            }
+        },
     ]
     collated = collate_temporal(batch)
     # Padded to length 3
@@ -168,23 +211,27 @@ def test_collate_temporal_variable_length():
 
 # ── 5. SinusoidalTimeEmbedding ────────────────────────────────────────────────
 
+
 def test_sinusoidal_time_embedding_shape():
-    from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding
+    from pyhealth.models.embedding import SinusoidalTimeEmbedding
+
     emb = SinusoidalTimeEmbedding(dim=64, max_hours=720.0)
-    t   = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]])  # (2, 3)
+    t = torch.tensor([[0.0, 12.0, 24.0], [0.0, 6.0, 48.0]])  # (2, 3)
     out = emb(t)
     assert out.shape == (2, 3, 64)
 
 
 def test_sinusoidal_different_times_differ():
-    from pyhealth.models.unified_embedding import SinusoidalTimeEmbedding
+    from pyhealth.models.embedding import SinusoidalTimeEmbedding
+
     emb = SinusoidalTimeEmbedding(dim=32)
-    t0  = emb(torch.tensor([0.0]))
-    t1  = emb(torch.tensor([24.0]))
+    t0 = emb(torch.tensor([0.0]))
+    t1 = emb(torch.tensor([24.0]))
     assert not torch.allclose(t0, t1)
 
 
-# ── 6. UnifiedMultimodalEmbeddingModel — code-only smoke test ─────────────────
+# ── 6. UnifiedMultimodalEmbeddingModel, code-only smoke test ─────────────────
+
 
 def _make_code_processors_and_inputs(batch_size=2, seq_len=5):
     """Build a minimal dataset mock with a single CODE-modality field."""
@@ -199,14 +246,16 @@ def _make_code_processors_and_inputs(batch_size=2, seq_len=5):
 
     # Fake batch dict (as produced by collate_temporal)
     value = torch.randint(1, vocab_size, (batch_size, seq_len))
-    time  = torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1)
+    time = (
+        torch.arange(seq_len, dtype=torch.float32).unsqueeze(0).expand(batch_size, -1)
+    )
     inputs = {"codes": {"value": value, "time": time}}
 
     return processors, inputs
 
 
 def test_unified_model_code_only():
-    from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
 
     processors, inputs = _make_code_processors_and_inputs()
     model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64)
@@ -214,8 +263,8 @@ def test_unified_model_code_only():
     out = model(inputs)
 
     assert "sequence" in out
-    assert "time"     in out
-    assert "mask"     in out
+    assert "time" in out
+    assert "mask" in out
 
     B, S, E = out["sequence"].shape
     assert B == 2
@@ -223,46 +272,484 @@ def test_unified_model_code_only():
     assert E == 64
 
 
+def test_unified_model_token_emb_is_content_only():
+    """``token_emb`` is the content embedding BEFORE time/type are added, in the
+    same temporally-sorted order as ``sequence``; i.e.
+    ``sequence == token_emb + time_embed(time) + type_embedding(type_ids)``."""
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+
+    processors, inputs = _make_code_processors_and_inputs()
+    model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=64)
+    model.eval()
+
+    out = model(inputs)
+
+    assert "token_emb" in out
+    assert out["token_emb"].shape == out["sequence"].shape
+
+    # Reconstruct sequence from the content-only token_emb + the position-derived
+    # time and type embeddings; they must match exactly.
+    recomposed = (
+        out["token_emb"]
+        + model.time_embed(out["time"])
+        + model.type_embedding(out["type_ids"])
+    )
+    assert torch.allclose(out["sequence"], recomposed, atol=1e-5)
+    # And token_emb is genuinely distinct from sequence (time/type were added).
+    assert not torch.allclose(out["token_emb"], out["sequence"])
+
+
 def test_unified_model_rejects_non_temporal():
-    from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
     from pyhealth.processors import SequenceProcessor
 
     bad_proc = SequenceProcessor()
-    with pytest.raises(TypeError, match="TemporalFeatureProcessor"):
-        UnifiedMultimodalEmbeddingModel(processors={"field": bad_proc}, embedding_dim=64)
+    try:
+        UnifiedMultimodalEmbeddingModel(
+            processors={"field": bad_proc},
+            embedding_dim=64,
+        )
+    except TypeError as exc:
+        assert "TemporalFeatureProcessor" in str(exc)
+    else:
+        raise AssertionError("Expected TypeError for non-temporal processor")
 
 
 def test_unified_model_gradient_flow():
     """Loss.backward() should propagate through time + type embeddings."""
-    from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
 
     processors, inputs = _make_code_processors_and_inputs()
     model = UnifiedMultimodalEmbeddingModel(processors=processors, embedding_dim=32)
 
-    out  = model(inputs)
+    out = model(inputs)
     loss = out["sequence"].mean()
     loss.backward()
 
     # type_embedding grad should be non-zero
     assert model.type_embedding.weight.grad is not None
-    assert model.time_embed.freqs.grad is None  # buffer, not parameter — OK
+    assert model.time_embed.freqs.grad is None  # buffer, not parameter, OK
 
 
 def test_unified_model_time_sort():
     """Events should be sorted by time ascending in the output."""
-    from pyhealth.models.unified_embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
     from pyhealth.processors import StageNetProcessor
 
     samples = [{"c": (None, ["a", "b"])}]
     proc = StageNetProcessor()
     proc.fit(samples, "c")
 
+    model = UnifiedMultimodalEmbeddingModel(processors={"c": proc}, embedding_dim=16)
+    # Reverse-order times
+    value = torch.tensor([[2, 1]])  # (1, 2)
+    time = torch.tensor([[10.0, 0.0]])  # t=10 then t=0 → should sort to [0, 10]
+    out = model({"c": {"value": value, "time": time}})
+    assert math.isclose(out["time"][0, 0].item(), 0.0, rel_tol=1e-6)
+    assert math.isclose(out["time"][0, 1].item(), 10.0, rel_tol=1e-6)
+
+
+# ── 7. field_embeddings: reuse pre-built unimodal encoder ─────────────────────
+
+
+def test_unified_field_embeddings_reuses_encoder():
+    """field_embeddings: encoder from a pre-built model is used in-place."""
+    import torch.nn as nn
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.processors import StageNetProcessor
+
+    proc = StageNetProcessor()
+    proc.fit([{"codes": (None, [f"c{i}" for i in range(5)])}], "codes")
+    vocab_size = proc.value_dim()
+
+    # Simulate a pre-built EmbeddingModel via a lightweight mock
+    pre_emb = nn.Embedding(vocab_size, 32)
+
+    class _MockEmbedModel:
+        embedding_dim = 32
+        embedding_layers = {"codes": pre_emb}
+
     model = UnifiedMultimodalEmbeddingModel(
-        processors={"c": proc}, embedding_dim=16
+        processors={"codes": proc},
+        embedding_dim=32,
+        field_embeddings={"codes": _MockEmbedModel()},
     )
-    # Reverse-order times
-    value = torch.tensor([[2, 1]])   # (1, 2)
-    time  = torch.tensor([[10.0, 0.0]])  # t=10 then t=0 → should sort to [0, 10]
-    out   = model({"c": {"value": value, "time": time}})
-    assert out["time"][0, 0].item() == pytest.approx(0.0)
-    assert out["time"][0, 1].item() == pytest.approx(10.0)
+    # The encoder registered should be the exact same object
+    assert model.encoders["codes"] is pre_emb
+
+
+def test_unified_field_embeddings_projection_added_on_dim_mismatch():
+    """When pre-built embedding_dim != unified embedding_dim, a projection is added."""
+    import torch.nn as nn
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.processors import StageNetProcessor
+
+    proc = StageNetProcessor()
+    proc.fit([{"codes": (None, ["A", "B"])}], "codes")
+    vocab_size = proc.value_dim()
+
+    pre_emb = nn.Embedding(vocab_size, 16)  # pre-built dim=16
+
+    class _MockEmbedModel:
+        embedding_dim = 16
+        embedding_layers = {"codes": pre_emb}
+
+    model = UnifiedMultimodalEmbeddingModel(
+        processors={"codes": proc},
+        embedding_dim=32,  # different from pre-built
+        field_embeddings={"codes": _MockEmbedModel()},
+    )
+    # A Sequential(pre_emb, nn.Linear(16→32)) should be built
+    assert isinstance(model.encoders["codes"], nn.Sequential)
+    # Forward should produce embedding_dim=32
+    value = torch.randint(1, vocab_size, (2, 3))
+    time = torch.arange(3, dtype=torch.float32).unsqueeze(0).expand(2, -1)
+    out = model({"codes": {"value": value, "time": time}})
+    assert out["sequence"].shape[-1] == 32
+
+
+def test_unified_field_embeddings_forward():
+    """End-to-end forward with field_embeddings reusing a CODE encoder."""
+    import torch.nn as nn
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.processors import StageNetProcessor
+
+    proc = StageNetProcessor()
+    proc.fit([{"codes": (None, [f"c{i}" for i in range(4)])}], "codes")
+    vocab_size = proc.value_dim()
+
+    pre_emb = nn.Embedding(vocab_size, 64)
+
+    class _MockEmbedModel:
+        embedding_dim = 64
+        embedding_layers = {"codes": pre_emb}
+
+    model = UnifiedMultimodalEmbeddingModel(
+        processors={"codes": proc},
+        embedding_dim=64,
+        field_embeddings={"codes": _MockEmbedModel()},
+    )
+    value = torch.randint(1, vocab_size, (3, 4))
+    time = torch.arange(4, dtype=torch.float32).unsqueeze(0).expand(3, -1)
+    out = model({"codes": {"value": value, "time": time}})
+
+    assert out["sequence"].shape == (3, 4, 64)
+    assert out["mask"].shape == (3, 4)
+
+
+def test_unified_text_encoder_shared_by_tokenizer():
+    """Token-based TEXT fields with the same tokenizer share one encoder."""
+    from types import SimpleNamespace
+    import torch.nn as nn
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.processors import ModalityType, TemporalFeatureProcessor
+
+    class _DummyTemporalTextProcessor(TemporalFeatureProcessor):
+        def __init__(self, tokenizer_model: str):
+            self.tokenizer_model = tokenizer_model
+
+        def process(self, value):
+            return value
+
+        def modality(self):
+            return ModalityType.TEXT
+
+        def value_dim(self):
+            return 0
+
+        def is_token(self):
+            return True
+
+        def schema(self):
+            return ("value", "mask", "time")
+
+        def dim(self):
+            return (2, 2, 1)
+
+        def spatial(self):
+            return (False, False)
+
+    class _DummyBert(nn.Module):
+        def __init__(self, hidden_size: int):
+            super().__init__()
+            self.config = SimpleNamespace(hidden_size=hidden_size)
+
+        def forward(self, input_ids=None, attention_mask=None):
+            if input_ids is None:
+                raise ValueError("input_ids is required")
+            b, l = input_ids.shape
+            hidden = self.config.hidden_size
+            out = torch.zeros(b, l, hidden)
+            return SimpleNamespace(last_hidden_state=out)
+
+    call_count = {"n": 0}
+
+    def _fake_from_pretrained(_name):
+        call_count["n"] += 1
+        return _DummyBert(hidden_size=48)
+
+    with patch("transformers.AutoModel.from_pretrained", _fake_from_pretrained):
+        processors = {
+            "discharge_note_times": _DummyTemporalTextProcessor("bert-base-uncased"),
+            "radiology_note_times": _DummyTemporalTextProcessor("bert-base-uncased"),
+        }
+
+        model = UnifiedMultimodalEmbeddingModel(
+            processors=processors,
+            embedding_dim=32,
+        )
+
+    # Both text fields reuse the same encoder instance.
+    assert (
+        model.encoders["discharge_note_times"] is model.encoders["radiology_note_times"]
+    )
+    assert call_count["n"] == 1
+
+    # Projections remain field-specific.
+    assert "discharge_note_times" in model.projections
+    assert "radiology_note_times" in model.projections
+    assert (
+        model.projections["discharge_note_times"]
+        is not model.projections["radiology_note_times"]
+    )
+
+
+# ── 8. Downstream models in unified mode ──────────────────────────────────────
+
+
+def _make_stagenet_dataset(n_codes: int = 5):
+    """Build a minimal SampleDataset with one StageNetProcessor field.
+
+    Time arrays are kept the same length as the code arrays so that the
+    temporal batch has consistent shapes.
+    """
+    from pyhealth.datasets import create_sample_dataset
+
+    codes_p0 = [f"c{i}" for i in range(n_codes)]
+    times_p0 = [float(i) for i in range(n_codes)]
+    codes_p1 = [f"c{i}" for i in range(2)]
+    times_p1 = [0.0, 1.0]
+
+    samples = [
+        {
+            "patient_id": "p0",
+            "visit_id": "v0",
+            "codes": (times_p0, codes_p0),
+            "label": 1,
+        },
+        {
+            "patient_id": "p1",
+            "visit_id": "v1",
+            "codes": (times_p1, codes_p1),
+            "label": 0,
+        },
+    ]
+    return create_sample_dataset(
+        samples,
+        input_schema={"codes": "stagenet"},
+        output_schema={"label": "binary"},
+        dataset_name="test_unified_downstream",
+    )
+
+
+def test_transformer_unified_mode():
+    """Transformer with unified_embedding uses a single backbone + forward works."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.transformer import Transformer
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = Transformer(dataset=dataset, embedding_dim=32, unified_embedding=unified)
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out and "logit" in out
+    out["loss"].backward()
+    # fc input size should be embedding_dim, not n_fields * embedding_dim
+    assert model.fc.in_features == 32
+
+
+def test_ehrmamba_unified_mode():
+    """EHRMamba with unified_embedding uses a single Mamba stack."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.ehrmamba import EHRMamba
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = EHRMamba(
+        dataset=dataset,
+        embedding_dim=32,
+        num_layers=1,
+        unified_embedding=unified,
+    )
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out
+    out["loss"].backward()
+    assert model.fc.in_features == 32
+
+
+def test_jamba_ehr_unified_mode():
+    """JambaEHR with unified_embedding uses a single JambaLayer."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.jamba_ehr import JambaEHR
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = JambaEHR(
+        dataset=dataset,
+        embedding_dim=32,
+        num_transformer_layers=1,
+        num_mamba_layers=1,
+        heads=2,
+        unified_embedding=unified,
+    )
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out
+    out["loss"].backward()
+    assert model.fc.in_features == 32
+
+
+def test_mlp_unified_mode():
+    """MLP with unified_embedding mean-pools the event sequence and produces valid output."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.mlp import MLP
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = MLP(
+        dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified
+    )
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out and "logit" in out
+    out["loss"].backward()
+    # fc input size should be hidden_dim, not n_fields * hidden_dim
+    assert model.fc.in_features == 32
+
+
+def test_rnn_unified_mode():
+    """RNN with unified_embedding uses a single RNN over the event sequence."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.rnn import RNN
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = RNN(
+        dataset=dataset, embedding_dim=32, hidden_dim=32, unified_embedding=unified
+    )
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out and "logit" in out
+    out["loss"].backward()
+    # fc input size should be hidden_dim, not n_fields * hidden_dim
+    assert model.fc.in_features == 32
+
+
+def test_bottleneck_transformer_unified_mode():
+    """BottleneckTransformer with unified_embedding uses n_modality=1 encoder and a single CLS token."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.embedding import UnifiedMultimodalEmbeddingModel
+    from pyhealth.models.bottleneck_transformer import BottleneckTransformer
+
+    dataset = _make_stagenet_dataset()
+    unified = UnifiedMultimodalEmbeddingModel(
+        processors=dataset.input_processors,
+        embedding_dim=32,
+    )
+    model = BottleneckTransformer(
+        dataset=dataset,
+        embedding_dim=32,
+        bottlenecks_n=2,
+        fusion_startidx=1,
+        num_layers=2,
+        heads=2,
+        unified_embedding=unified,
+    )
+
+    loader = get_dataloader(dataset, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+
+    assert "loss" in out and "y_prob" in out and "logit" in out
+    out["loss"].backward()
+    # fc input must be embedding_dim (CLS token), not n_fields * embedding_dim
+    assert model.fc.in_features == 32
+    # Encoder must be configured for n_modality=1
+    assert model.encoder.n_modality == 1
+
+
+def test_unified_per_field_backward_compat():
+    """Models without unified_embedding still work in per-field mode."""
+    from pyhealth.datasets.utils import get_dataloader
+    from pyhealth.models.transformer import Transformer
+
+    dataset = _make_stagenet_dataset()
+    # Uses SequenceProcessor-style input_schema for per-field mode
+    from pyhealth.datasets import create_sample_dataset
+
+    samples = [
+        {"patient_id": "p0", "visit_id": "v0", "codes": ["A", "B", "C"], "label": 1},
+        {"patient_id": "p1", "visit_id": "v1", "codes": ["D", "E"], "label": 0},
+    ]
+    ds = create_sample_dataset(
+        samples,
+        input_schema={"codes": "sequence"},
+        output_schema={"label": "binary"},
+        dataset_name="test_compat",
+    )
+    model = Transformer(dataset=ds, embedding_dim=32)
+    loader = get_dataloader(ds, batch_size=2, shuffle=False)
+    batch = next(iter(loader))
+    out = model(**batch)
+    assert "loss" in out
+    out["loss"].backward()
+
+
+def load_tests(loader, tests, pattern):
+    """Expose top-level test_ functions to unittest discovery."""
+    suite = unittest.TestSuite()
+    namespace = globals()
+    for name in sorted(namespace):
+        if name.startswith("test_") and callable(namespace[name]):
+            suite.addTest(unittest.FunctionTestCase(namespace[name]))
+    return suite
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/tools/summarize_sweep.py b/tools/summarize_sweep.py
new file mode 100644
index 000000000..07100e2d8
--- /dev/null
+++ b/tools/summarize_sweep.py
@@ -0,0 +1,241 @@
+"""Summarize a completed E2E sweep.
+
+Reads metrics_history.json and predictions_*.csv from each experiment directory
+and prints a results table + saves plots.
+
+Usage:
+    python tools/summarize_sweep.py --sweep-dir ~/Downloads/e2e_sweep
+    python tools/summarize_sweep.py --sweep-dir ~/Downloads/e2e_sweep --plot
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+
+def load_metrics(exp_dir: Path) -> dict | None:
+    p = exp_dir / "metrics_history.json"
+    if not p.exists():
+        return None
+    with p.open() as f:
+        return json.load(f)
+
+
+def load_log_metrics(exp_dir: Path) -> dict | None:
+    """Fallback: parse log.txt for best val metrics when metrics_history.json absent."""
+    log = exp_dir / "log.txt"
+    if not log.exists():
+        return None
+
+    best = {}
+    current_epoch = {}
+    for raw in log.read_text().splitlines():
+        # strip optional "YYYY-MM-DD HH:MM:SS " timestamp prefix
+        line = raw.strip()
+        parts = line.split(" ", 2)
+        if len(parts) == 3 and len(parts[0]) == 10 and parts[0][4] == "-":
+            line = parts[2].strip()
+        # detect epoch boundary
+        if "--- Eval epoch-" in line:
+            current_epoch = {}
+        for metric in ("pr_auc", "roc_auc", "accuracy"):
+            if line.startswith(metric + ":"):
+                try:
+                    current_epoch[metric] = float(line.split(":", 1)[1].strip())
+                except ValueError:
+                    pass
+        if "New best pr_auc" in line:
+            best = dict(current_epoch)
+    return best if best else None
+
+
+def parse_exp_name(name: str) -> tuple[str, int]:
+    # e.g. "bottleneck_transformer_seed42"
+    parts = name.rsplit("_seed", 1)
+    return parts[0], int(parts[1])
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--sweep-dir", type=str, required=True)
+    parser.add_argument("--plot", action="store_true", help="Save convergence plots")
+    args = parser.parse_args()
+
+    sweep_dir = Path(args.sweep_dir).expanduser()
+    if not sweep_dir.exists():
+        raise SystemExit(f"Directory not found: {sweep_dir}")
+
+    rows = []
+    histories = {}
+
+    for exp_dir in sorted(sweep_dir.iterdir()):
+        if not exp_dir.is_dir():
+            continue
+        try:
+            model, seed = parse_exp_name(exp_dir.name)
+        except (ValueError, IndexError):
+            continue
+
+        hist = load_metrics(exp_dir)
+        if hist:
+            # hist is a flat list of epoch dicts with keys like val_pr_auc, train_loss, etc.
+            valid = [e for e in hist if e.get("val_pr_auc") is not None]
+            if valid:
+                best_e = max(valid, key=lambda e: e.get("val_pr_auc", -1))
+                best = {
+                    "pr_auc": best_e.get("val_pr_auc", float("nan")),
+                    "roc_auc": best_e.get("val_roc_auc", float("nan")),
+                    "accuracy": best_e.get("val_accuracy", float("nan")),
+                }
+                nan_epochs = [e["epoch"] for e in hist if np.isnan(e.get("train_loss", 0))]
+            else:
+                best, nan_epochs = {}, []
+            histories[exp_dir.name] = hist
+        else:
+            # fallback to log parsing
+            best = load_log_metrics(exp_dir) or {}
+            nan_epochs = []
+            # check log for nan
+            log = exp_dir / "log.txt"
+            if log.exists():
+                nan_epochs = [i for i, l in enumerate(log.read_text().splitlines()) if "loss: nan" in l]
+
+        rows.append({
+            "model": model,
+            "seed": seed,
+            "pr_auc": best.get("pr_auc", float("nan")),
+            "roc_auc": best.get("roc_auc", float("nan")),
+            "accuracy": best.get("accuracy", float("nan")),
+            "nan_epochs": len(nan_epochs),
+            "status": "NaN" if nan_epochs else ("OK" if best else "missing"),
+        })
+
+    if not rows:
+        raise SystemExit("No experiment directories found.")
+
+    # ---- Print summary table ------------------------------------
+    preferred_order = [
+        "mlp",
+        "rnn",
+        "transformer",
+        "bottleneck_transformer",
+        "ehrmamba",
+        "jambaehr",
+    ]
+    seen_models = sorted({r["model"] for r in rows})
+    MODELS = [m for m in preferred_order if m in seen_models] + [
+        m for m in seen_models if m not in preferred_order
+    ]
+    SEEDS = sorted({r["seed"] for r in rows})
+
+    row_map = {(r["model"], r["seed"]): r for r in rows}
+
+    col_w = 28
+    seed_w = 32
+
+    header = f"{'Model':<{col_w}}" + "".join(f"{'seed' + str(s):<{seed_w}}" for s in SEEDS) + f"{'Mean PR-AUC':>12}"
+    print("\n" + "=" * len(header))
+    print("SWEEP RESULTS — Best Val PR-AUC (roc_auc / accuracy)")
+    print("=" * len(header))
+    print(header)
+    print("-" * len(header))
+
+    for model in MODELS:
+        model_rows = [row_map.get((model, s)) for s in SEEDS]
+        pr_aucs = [r["pr_auc"] for r in model_rows if r and not np.isnan(r["pr_auc"])]
+        mean_pr = np.mean(pr_aucs) if pr_aucs else float("nan")
+
+        cells = []
+        for r in model_rows:
+            if r is None:
+                cells.append(f"{'—':<{seed_w}}")
+            elif r["status"] == "missing" or np.isnan(r["pr_auc"]):
+                cells.append(f"{'missing':<{seed_w}}")
+            else:
+                marker = "* " if r["status"] == "NaN" else ""
+                cell = f"{marker}{r['pr_auc']:.4f} ({r['roc_auc']:.4f}/{r['accuracy']:.4f})"
+                cells.append(f"{cell:<{seed_w}}")
+
+        mean_str = f"{mean_pr:.4f}" if not np.isnan(mean_pr) else "—"
+        print(f"{model:<{col_w}}" + "".join(cells) + f"{mean_str:>12}")
+
+    print("=" * len(header))
+
+    # ---- Per-model mean ± std -----------------------------------
+    print("\nMean ± Std PR-AUC across seeds:")
+    for model in MODELS:
+        vals = [row_map[(model, s)]["pr_auc"] for s in SEEDS
+                if (model, s) in row_map and not np.isnan(row_map[(model, s)]["pr_auc"])]
+        if vals:
+            print(f"  {model:<35} {np.mean(vals):.4f} ± {np.std(vals):.4f}  (n={len(vals)})")
+        else:
+            print(f"  {model:<35} no valid runs")
+
+    # ---- NaN warnings -------------------------------------------
+    nan_runs = [r for r in rows if r["status"] == "NaN"]
+    if nan_runs:
+        print(f"\n⚠  NaN runs detected ({len(nan_runs)}):")
+        for r in nan_runs:
+            print(f"   {r['model']}_seed{r['seed']}")
+
+    # ---- Plots --------------------------------------------------
+    if args.plot:
+        try:
+            import matplotlib.pyplot as plt
+        except ImportError:
+            print("\nmatplotlib not installed — skipping plots (pip install matplotlib)")
+            return
+
+        plot_dir = sweep_dir / "plots"
+        plot_dir.mkdir(exist_ok=True)
+
+        for model in MODELS:
+            fig, ax = plt.subplots(figsize=(8, 4))
+            plotted = False
+            for seed in SEEDS:
+                key = f"{model}_seed{seed}"
+                hist = histories.get(key)
+                if not hist:
+                    continue
+                pr_aucs = [e.get("val_pr_auc", float("nan")) for e in hist if "val_pr_auc" in e]
+                ax.plot(pr_aucs, label=f"seed {seed}")
+                plotted = True
+            if plotted:
+                ax.set_title(f"{model} — Val PR-AUC over epochs")
+                ax.set_xlabel("Epoch")
+                ax.set_ylabel("PR-AUC")
+                ax.legend()
+                ax.grid(True, alpha=0.3)
+                out = plot_dir / f"{model}_pr_auc.png"
+                fig.savefig(out, dpi=150, bbox_inches="tight")
+                print(f"Saved {out}")
+            plt.close(fig)
+
+        # Combined final-epoch bar chart
+        fig, ax = plt.subplots(figsize=(10, 5))
+        x = np.arange(len(MODELS))
+        width = 0.25
+        for i, seed in enumerate(SEEDS):
+            vals = []
+            for model in MODELS:
+                r = row_map.get((model, seed))
+                vals.append(r["pr_auc"] if r and not np.isnan(r["pr_auc"]) else 0)
+            ax.bar(x + i * width, vals, width, label=f"seed {seed}")
+        ax.set_xticks(x + width)
+        ax.set_xticklabels(MODELS, rotation=15, ha="right")
+        ax.set_ylabel("Best Val PR-AUC")
+        ax.set_title("Model Comparison — Best Val PR-AUC")
+        ax.legend()
+        ax.grid(True, alpha=0.3, axis="y")
+        out = plot_dir / "model_comparison.png"
+        fig.savefig(out, dpi=150, bbox_inches="tight")
+        print(f"Saved {out}")
+        plt.close(fig)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/wandb/latest-run b/wandb/latest-run
new file mode 120000
index 000000000..86a367304
--- /dev/null
+++ b/wandb/latest-run
@@ -0,0 +1 @@
+run-20260807_042507-cmlmjn78
\ No newline at end of file
diff --git a/wandb/run-20260705_212337-99vdbgbw/files/config.yaml b/wandb/run-20260705_212337-99vdbgbw/files/config.yaml
new file mode 100644
index 000000000..6e90eecaf
--- /dev/null
+++ b/wandb/run-20260705_212337-99vdbgbw/files/config.yaml
@@ -0,0 +1,177 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            hml53lcj4oy2dckmncuhvw4ipv9nku1x:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --dev
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "64"
+                    - --hidden-dim
+                    - "64"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15924445184"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-06T02:23:37.858585Z"
+                writerId: hml53lcj4oy2dckmncuhvw4ipv9nku1x
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+dev_mode:
+    value: true
+exp_name:
+    value: rnn_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: rnn
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260705_212337-99vdbgbw/files/requirements.txt b/wandb/run-20260705_212337-99vdbgbw/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260705_212337-99vdbgbw/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260705_212337-99vdbgbw/files/wandb-metadata.json b/wandb/run-20260705_212337-99vdbgbw/files/wandb-metadata.json
new file mode 100644
index 000000000..e46394335
--- /dev/null
+++ b/wandb/run-20260705_212337-99vdbgbw/files/wandb-metadata.json
@@ -0,0 +1,128 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-06T02:23:37.858585Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--dev",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "64",
+    "--hidden-dim",
+    "64",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "1",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15924445184"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "hml53lcj4oy2dckmncuhvw4ipv9nku1x"
+}
\ No newline at end of file
diff --git a/wandb/run-20260705_212337-99vdbgbw/files/wandb-summary.json b/wandb/run-20260705_212337-99vdbgbw/files/wandb-summary.json
new file mode 100644
index 000000000..c03270709
--- /dev/null
+++ b/wandb/run-20260705_212337-99vdbgbw/files/wandb-summary.json
@@ -0,0 +1 @@
+{"val_loss":0.21904340386390686,"_step":8,"global_step":171,"_timestamp":1.7833046309801621e+09,"epoch":8,"train_vram_allocated_mb":16.5615234375,"train_loss":0.16386246249863975,"_runtime":13.321371775,"val_pr_auc":0.6118524332810047,"_wandb":{"runtime":13},"val_accuracy":0.9054054054054054,"val_roc_auc":0.9189765458422174,"epoch_time_s":1.031,"val_f1":0,"train_vram_peak_mb":89.7705078125}
\ No newline at end of file
diff --git a/wandb/run-20260705_212337-99vdbgbw/run-99vdbgbw.wandb b/wandb/run-20260705_212337-99vdbgbw/run-99vdbgbw.wandb
new file mode 100644
index 000000000..6c8e9d222
Binary files /dev/null and b/wandb/run-20260705_212337-99vdbgbw/run-99vdbgbw.wandb differ
diff --git a/wandb/run-20260705_225706-wkdkg5al/files/config.yaml b/wandb/run-20260705_225706-wkdkg5al/files/config.yaml
new file mode 100644
index 000000000..f7d63db5a
--- /dev/null
+++ b/wandb/run-20260705_225706-wkdkg5al/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            78phm1uydmhscrrd5pm56g60304gckf2:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "64"
+                    - --hidden-dim
+                    - "64"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15926824960"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-06T03:57:06.652653Z"
+                writerId: 78phm1uydmhscrrd5pm56g60304gckf2
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137689
+data/n_neg_val:
+    value: 17245
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6897
+data/n_pos_val:
+    value: 828
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.047701713858879835
+data/pos_rate_val:
+    value: 0.04581419797487966
+dev_mode:
+    value: false
+exp_name:
+    value: rnn_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: rnn
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260705_225706-wkdkg5al/files/requirements.txt b/wandb/run-20260705_225706-wkdkg5al/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260705_225706-wkdkg5al/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260705_225706-wkdkg5al/files/wandb-metadata.json b/wandb/run-20260705_225706-wkdkg5al/files/wandb-metadata.json
new file mode 100644
index 000000000..2a3b82e3b
--- /dev/null
+++ b/wandb/run-20260705_225706-wkdkg5al/files/wandb-metadata.json
@@ -0,0 +1,127 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-06T03:57:06.652653Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "64",
+    "--hidden-dim",
+    "64",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "1",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15926824960"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "78phm1uydmhscrrd5pm56g60304gckf2"
+}
\ No newline at end of file
diff --git a/wandb/run-20260705_225706-wkdkg5al/files/wandb-summary.json b/wandb/run-20260705_225706-wkdkg5al/files/wandb-summary.json
new file mode 100644
index 000000000..f02a49475
--- /dev/null
+++ b/wandb/run-20260705_225706-wkdkg5al/files/wandb-summary.json
@@ -0,0 +1 @@
+{"epoch_time_s":110.619,"val_loss":0.10411418034835199,"val_f1":0.556060606060606,"_timestamp":1.7833122336183228e+09,"val_accuracy":0.9675759420129475,"val_pr_auc":0.6126895956406891,"epoch":16,"val_roc_auc":0.9231699519429422,"train_vram_allocated_mb":16.5615234375,"_step":16,"train_vram_peak_mb":337.224609375,"_runtime":2014.862620307,"train_loss":0.1075635792495459,"_wandb":{"runtime":2014},"global_step":76823}
\ No newline at end of file
diff --git a/wandb/run-20260705_225706-wkdkg5al/run-wkdkg5al.wandb b/wandb/run-20260705_225706-wkdkg5al/run-wkdkg5al.wandb
new file mode 100644
index 000000000..de6d93ad0
Binary files /dev/null and b/wandb/run-20260705_225706-wkdkg5al/run-wkdkg5al.wandb differ
diff --git a/wandb/run-20260706_003500-qtt8p5nt/files/config.yaml b/wandb/run-20260706_003500-qtt8p5nt/files/config.yaml
new file mode 100644
index 000000000..5344e093e
--- /dev/null
+++ b/wandb/run-20260706_003500-qtt8p5nt/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            j0yp0oup48o4y5ygipytp4uly9uwz9s7:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - transformer
+                    - --embedding-dim
+                    - "64"
+                    - --hidden-dim
+                    - "64"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15926120448"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-06T05:35:00.005931Z"
+                writerId: j0yp0oup48o4y5ygipytp4uly9uwz9s7
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137687
+data/n_neg_val:
+    value: 17247
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6899
+data/n_pos_val:
+    value: 826
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.04771554645678005
+data/pos_rate_val:
+    value: 0.04570353566093067
+dev_mode:
+    value: false
+exp_name:
+    value: transformer_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: transformer
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260706_003500-qtt8p5nt/files/requirements.txt b/wandb/run-20260706_003500-qtt8p5nt/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260706_003500-qtt8p5nt/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260706_003500-qtt8p5nt/files/wandb-metadata.json b/wandb/run-20260706_003500-qtt8p5nt/files/wandb-metadata.json
new file mode 100644
index 000000000..4e45357e6
--- /dev/null
+++ b/wandb/run-20260706_003500-qtt8p5nt/files/wandb-metadata.json
@@ -0,0 +1,127 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-06T05:35:00.005931Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "transformer",
+    "--embedding-dim",
+    "64",
+    "--hidden-dim",
+    "64",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15926120448"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "j0yp0oup48o4y5ygipytp4uly9uwz9s7"
+}
\ No newline at end of file
diff --git a/wandb/run-20260706_003500-qtt8p5nt/files/wandb-summary.json b/wandb/run-20260706_003500-qtt8p5nt/files/wandb-summary.json
new file mode 100644
index 000000000..b9f0b4ec6
--- /dev/null
+++ b/wandb/run-20260706_003500-qtt8p5nt/files/wandb-summary.json
@@ -0,0 +1 @@
+{"val_f1":0.5285379202501954,"_runtime":2068.907608139,"_step":15,"train_loss":0.11486118818489201,"val_pr_auc":0.587121664919961,"val_roc_auc":0.9194649215058069,"val_loss":0.11249546522884506,"train_vram_allocated_mb":17.42041015625,"epoch_time_s":120.947,"val_accuracy":0.9666353123443812,"_wandb":{"runtime":2068},"_timestamp":1.7833181601982453e+09,"epoch":15,"global_step":72304,"train_vram_peak_mb":27835.92724609375}
\ No newline at end of file
diff --git a/wandb/run-20260706_003500-qtt8p5nt/run-qtt8p5nt.wandb b/wandb/run-20260706_003500-qtt8p5nt/run-qtt8p5nt.wandb
new file mode 100644
index 000000000..3aee362f9
Binary files /dev/null and b/wandb/run-20260706_003500-qtt8p5nt/run-qtt8p5nt.wandb differ
diff --git a/wandb/run-20260709_075941-6tix8xfj/files/config.yaml b/wandb/run-20260709_075941-6tix8xfj/files/config.yaml
new file mode 100644
index 000000000..8cef2fb92
--- /dev/null
+++ b/wandb/run-20260709_075941-6tix8xfj/files/config.yaml
@@ -0,0 +1,214 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            qckf2oaltjxv63411k18pz5dgotnwnp3:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - bottleneck_transformer
+                    - --embedding-dim
+                    - "64"
+                    - --hidden-dim
+                    - "64"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --bottlenecks-n
+                    - "4"
+                    - --fusion-startidx
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --max-grad-norm
+                    - "0.5"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15992446976"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-09T12:59:41.103820Z"
+                writerId: qckf2oaltjxv63411k18pz5dgotnwnp3
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137687
+data/n_neg_val:
+    value: 17247
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6899
+data/n_pos_val:
+    value: 826
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.04771554645678005
+data/pos_rate_val:
+    value: 0.04570353566093067
+dev_mode:
+    value: false
+exp_name:
+    value: bottleneck_transformer_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: bottleneck_transformer
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260709_075941-6tix8xfj/files/requirements.txt b/wandb/run-20260709_075941-6tix8xfj/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260709_075941-6tix8xfj/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260709_075941-6tix8xfj/files/wandb-metadata.json b/wandb/run-20260709_075941-6tix8xfj/files/wandb-metadata.json
new file mode 100644
index 000000000..15261af46
--- /dev/null
+++ b/wandb/run-20260709_075941-6tix8xfj/files/wandb-metadata.json
@@ -0,0 +1,133 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-09T12:59:41.103820Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "bottleneck_transformer",
+    "--embedding-dim",
+    "64",
+    "--hidden-dim",
+    "64",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--bottlenecks-n",
+    "4",
+    "--fusion-startidx",
+    "1",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--max-grad-norm",
+    "0.5",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15992446976"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "qckf2oaltjxv63411k18pz5dgotnwnp3"
+}
\ No newline at end of file
diff --git a/wandb/run-20260709_075941-6tix8xfj/files/wandb-summary.json b/wandb/run-20260709_075941-6tix8xfj/files/wandb-summary.json
new file mode 100644
index 000000000..a6d5755ee
--- /dev/null
+++ b/wandb/run-20260709_075941-6tix8xfj/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_wandb":{"runtime":941},"val_accuracy":0.9568970287168704,"global_step":36152,"_timestamp":1.7836029155096333e+09,"epoch_time_s":113.818,"val_pr_auc":0.3816664877599453,"_step":7,"train_vram_allocated_mb":18.37109375,"val_roc_auc":0.8874258020940864,"epoch":7,"_runtime":941.223895169,"val_f1":0.14676889375684557,"val_loss":0.13282482986363162,"train_vram_peak_mb":1034.2265625,"train_loss":0.13406056100859146}
\ No newline at end of file
diff --git a/wandb/run-20260709_075941-6tix8xfj/run-6tix8xfj.wandb b/wandb/run-20260709_075941-6tix8xfj/run-6tix8xfj.wandb
new file mode 100644
index 000000000..0366fc358
Binary files /dev/null and b/wandb/run-20260709_075941-6tix8xfj/run-6tix8xfj.wandb differ
diff --git a/wandb/run-20260709_084129-zwavnt2y/files/config.yaml b/wandb/run-20260709_084129-zwavnt2y/files/config.yaml
new file mode 100644
index 000000000..b223a3ad9
--- /dev/null
+++ b/wandb/run-20260709_084129-zwavnt2y/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            23wt9flz2nbfkcb03ru6ayy0fq87jz7l:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - ehrmamba
+                    - --embedding-dim
+                    - "64"
+                    - --num-layers
+                    - "2"
+                    - --mamba-state-size
+                    - "16"
+                    - --mamba-conv-kernel
+                    - "4"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15992516608"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-09T13:41:29.980901Z"
+                writerId: 23wt9flz2nbfkcb03ru6ayy0fq87jz7l
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137687
+data/n_neg_val:
+    value: 17247
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6899
+data/n_pos_val:
+    value: 826
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.04771554645678005
+data/pos_rate_val:
+    value: 0.04570353566093067
+dev_mode:
+    value: false
+exp_name:
+    value: ehrmamba_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: ehrmamba
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260709_084129-zwavnt2y/files/requirements.txt b/wandb/run-20260709_084129-zwavnt2y/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260709_084129-zwavnt2y/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260709_084129-zwavnt2y/files/wandb-metadata.json b/wandb/run-20260709_084129-zwavnt2y/files/wandb-metadata.json
new file mode 100644
index 000000000..ccbf644bd
--- /dev/null
+++ b/wandb/run-20260709_084129-zwavnt2y/files/wandb-metadata.json
@@ -0,0 +1,127 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-09T13:41:29.980901Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "ehrmamba",
+    "--embedding-dim",
+    "64",
+    "--num-layers",
+    "2",
+    "--mamba-state-size",
+    "16",
+    "--mamba-conv-kernel",
+    "4",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15992516608"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "23wt9flz2nbfkcb03ru6ayy0fq87jz7l"
+}
\ No newline at end of file
diff --git a/wandb/run-20260709_084129-zwavnt2y/files/wandb-summary.json b/wandb/run-20260709_084129-zwavnt2y/files/wandb-summary.json
new file mode 100644
index 000000000..576edb04e
--- /dev/null
+++ b/wandb/run-20260709_084129-zwavnt2y/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_step":28,"epoch":28,"_timestamp":1.7836081585907967e+09,"_wandb":{"runtime":3674},"_runtime":3674.76575785,"val_loss":0.10117696542787341,"train_vram_peak_mb":1201.97607421875,"val_roc_auc":0.9249758634375267,"train_loss":0.10581304661364507,"val_f1":0.5106022052586938,"epoch_time_s":117.448,"val_accuracy":0.9680739224257179,"train_vram_allocated_mb":18.49853515625,"global_step":131051,"val_pr_auc":0.6226378916971728}
\ No newline at end of file
diff --git a/wandb/run-20260709_084129-zwavnt2y/run-zwavnt2y.wandb b/wandb/run-20260709_084129-zwavnt2y/run-zwavnt2y.wandb
new file mode 100644
index 000000000..923514f65
Binary files /dev/null and b/wandb/run-20260709_084129-zwavnt2y/run-zwavnt2y.wandb differ
diff --git a/wandb/run-20260709_100158-shou8a8s/files/config.yaml b/wandb/run-20260709_100158-shou8a8s/files/config.yaml
new file mode 100644
index 000000000..37107aeb2
--- /dev/null
+++ b/wandb/run-20260709_100158-shou8a8s/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            6lquc348m0jilcaa1sixo3l7k7b94lks:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - ehrmamba
+                    - --embedding-dim
+                    - "64"
+                    - --num-layers
+                    - "2"
+                    - --mamba-state-size
+                    - "16"
+                    - --mamba-conv-kernel
+                    - "4"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15993667584"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-09T15:01:58.665622Z"
+                writerId: 6lquc348m0jilcaa1sixo3l7k7b94lks
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137687
+data/n_neg_val:
+    value: 17247
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6899
+data/n_pos_val:
+    value: 826
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.04771554645678005
+data/pos_rate_val:
+    value: 0.04570353566093067
+dev_mode:
+    value: false
+exp_name:
+    value: ehrmamba_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: ehrmamba
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260709_100158-shou8a8s/files/requirements.txt b/wandb/run-20260709_100158-shou8a8s/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260709_100158-shou8a8s/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260709_100158-shou8a8s/files/wandb-metadata.json b/wandb/run-20260709_100158-shou8a8s/files/wandb-metadata.json
new file mode 100644
index 000000000..a048e6359
--- /dev/null
+++ b/wandb/run-20260709_100158-shou8a8s/files/wandb-metadata.json
@@ -0,0 +1,127 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-09T15:01:58.665622Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "ehrmamba",
+    "--embedding-dim",
+    "64",
+    "--num-layers",
+    "2",
+    "--mamba-state-size",
+    "16",
+    "--mamba-conv-kernel",
+    "4",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15993667584"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "6lquc348m0jilcaa1sixo3l7k7b94lks"
+}
\ No newline at end of file
diff --git a/wandb/run-20260709_100158-shou8a8s/files/wandb-summary.json b/wandb/run-20260709_100158-shou8a8s/files/wandb-summary.json
new file mode 100644
index 000000000..7be092d4e
--- /dev/null
+++ b/wandb/run-20260709_100158-shou8a8s/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_runtime":3610.742152168,"train_vram_allocated_mb":18.49853515625,"val_accuracy":0.9680739224257179,"val_pr_auc":0.6226378916971728,"train_vram_peak_mb":1201.97607421875,"_step":28,"epoch":28,"epoch_time_s":116.148,"global_step":131051,"val_f1":0.5106022052586938,"_timestamp":1.7836129221090653e+09,"val_loss":0.10117696542787341,"val_roc_auc":0.9249758634375267,"_wandb":{"runtime":3610},"train_loss":0.10581304661364507}
\ No newline at end of file
diff --git a/wandb/run-20260709_100158-shou8a8s/run-shou8a8s.wandb b/wandb/run-20260709_100158-shou8a8s/run-shou8a8s.wandb
new file mode 100644
index 000000000..9a44e7bbf
Binary files /dev/null and b/wandb/run-20260709_100158-shou8a8s/run-shou8a8s.wandb differ
diff --git a/wandb/run-20260709_121536-o96vhptf/files/config.yaml b/wandb/run-20260709_121536-o96vhptf/files/config.yaml
new file mode 100644
index 000000000..1f4cddbef
--- /dev/null
+++ b/wandb/run-20260709_121536-o96vhptf/files/config.yaml
@@ -0,0 +1,212 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            029m4i5aoeymxrchhou5j0mvz2rz22qo:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - jambaehr
+                    - --embedding-dim
+                    - "64"
+                    - --heads
+                    - "4"
+                    - --jamba-transformer-layers
+                    - "2"
+                    - --jamba-mamba-layers
+                    - "6"
+                    - --mamba-state-size
+                    - "16"
+                    - --mamba-conv-kernel
+                    - "4"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15995015168"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 10fabcefeadfd7a87586e898f5250aff0166aed2
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-09T17:15:36.083502Z"
+                writerId: 029m4i5aoeymxrchhou5j0mvz2rz22qo
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+arch/embedding_dim:
+    value: 64
+arch/hidden_dim:
+    value: 64
+arch/rnn_layers:
+    value: 1
+arch/rnn_type:
+    value: GRU
+data/cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+data/n_neg_test:
+    value: 17197
+data/n_neg_total:
+    value: 172131
+data/n_neg_train:
+    value: 137687
+data/n_neg_val:
+    value: 17247
+data/n_pos_test:
+    value: 877
+data/n_pos_total:
+    value: 8602
+data/n_pos_train:
+    value: 6899
+data/n_pos_val:
+    value: 826
+data/n_test:
+    value: 18074
+data/n_total:
+    value: 180733
+data/n_train:
+    value: 144586
+data/n_val:
+    value: 18073
+data/pos_rate_test:
+    value: 0.04852273984729446
+data/pos_rate_total:
+    value: 0.047595071182351865
+data/pos_rate_train:
+    value: 0.04771554645678005
+data/pos_rate_val:
+    value: 0.04570353566093067
+dev_mode:
+    value: false
+exp_name:
+    value: jambaehr_seed12
+hp/batch_size:
+    value: 32
+hp/dropout:
+    value: 0.1
+hp/epochs:
+    value: 50
+hp/lr:
+    value: 0.001
+hp/weight_decay:
+    value: 1e-05
+model:
+    value: jambaehr
+paths/ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+paths/output_dir:
+    value: /home/wp14/output
+seed:
+    value: 12
+task:
+    value: labs
diff --git a/wandb/run-20260709_121536-o96vhptf/files/requirements.txt b/wandb/run-20260709_121536-o96vhptf/files/requirements.txt
new file mode 100644
index 000000000..05fdbfec4
--- /dev/null
+++ b/wandb/run-20260709_121536-o96vhptf/files/requirements.txt
@@ -0,0 +1,221 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
diff --git a/wandb/run-20260709_121536-o96vhptf/files/wandb-metadata.json b/wandb/run-20260709_121536-o96vhptf/files/wandb-metadata.json
new file mode 100644
index 000000000..5df9eb7bd
--- /dev/null
+++ b/wandb/run-20260709_121536-o96vhptf/files/wandb-metadata.json
@@ -0,0 +1,131 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-09T17:15:36.083502Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "jambaehr",
+    "--embedding-dim",
+    "64",
+    "--heads",
+    "4",
+    "--jamba-transformer-layers",
+    "2",
+    "--jamba-mamba-layers",
+    "6",
+    "--mamba-state-size",
+    "16",
+    "--mamba-conv-kernel",
+    "4",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "10fabcefeadfd7a87586e898f5250aff0166aed2"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15995015168"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "029m4i5aoeymxrchhou5j0mvz2rz22qo"
+}
\ No newline at end of file
diff --git a/wandb/run-20260709_121536-o96vhptf/files/wandb-summary.json b/wandb/run-20260709_121536-o96vhptf/files/wandb-summary.json
new file mode 100644
index 000000000..6aaa62846
--- /dev/null
+++ b/wandb/run-20260709_121536-o96vhptf/files/wandb-summary.json
@@ -0,0 +1 @@
+{"train_vram_peak_mb":30905.638671875,"global_step":76823,"train_vram_allocated_mb":24.09423828125,"val_roc_auc":0.9268825711486337,"val_pr_auc":0.6154380670463732,"val_accuracy":0.9664139877164831,"train_loss":0.10747217934584427,"epoch_time_s":230.618,"val_loss":0.10395209964166964,"_wandb":{"runtime":4134},"_timestamp":1.7836214602059822e+09,"_step":16,"_runtime":4134.879882944,"val_f1":0.5716302046577276,"epoch":16}
\ No newline at end of file
diff --git a/wandb/run-20260709_121536-o96vhptf/run-o96vhptf.wandb b/wandb/run-20260709_121536-o96vhptf/run-o96vhptf.wandb
new file mode 100644
index 000000000..a87c8c5d1
Binary files /dev/null and b/wandb/run-20260709_121536-o96vhptf/run-o96vhptf.wandb differ
diff --git a/wandb/run-20260714_234904-mj2wwse4/files/config.yaml b/wandb/run-20260714_234904-mj2wwse4/files/config.yaml
new file mode 100644
index 000000000..109321f23
--- /dev/null
+++ b/wandb/run-20260714_234904-mj2wwse4/files/config.yaml
@@ -0,0 +1,228 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            tt0qgyvtwsznu5isg930c9z2wtm7ruef:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-only
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15968632832"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-15T04:49:04.128965Z"
+                writerId: tt0qgyvtwsznu5isg930c9z2wtm7ruef
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 16
+                - 61
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: null
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: null
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-only
+wandb_run_name:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260714_234904-mj2wwse4/files/requirements.txt b/wandb/run-20260714_234904-mj2wwse4/files/requirements.txt
new file mode 100644
index 000000000..d84484f0e
--- /dev/null
+++ b/wandb/run-20260714_234904-mj2wwse4/files/requirements.txt
@@ -0,0 +1,224 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
diff --git a/wandb/run-20260714_234904-mj2wwse4/files/wandb-metadata.json b/wandb/run-20260714_234904-mj2wwse4/files/wandb-metadata.json
new file mode 100644
index 000000000..e7a4c113d
--- /dev/null
+++ b/wandb/run-20260714_234904-mj2wwse4/files/wandb-metadata.json
@@ -0,0 +1,128 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-15T04:49:04.128965Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-only",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15968632832"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "tt0qgyvtwsznu5isg930c9z2wtm7ruef"
+}
\ No newline at end of file
diff --git a/wandb/run-20260714_234904-mj2wwse4/files/wandb-summary.json b/wandb/run-20260714_234904-mj2wwse4/files/wandb-summary.json
new file mode 100644
index 000000000..cf2786445
--- /dev/null
+++ b/wandb/run-20260714_234904-mj2wwse4/files/wandb-summary.json
@@ -0,0 +1 @@
+{"epoch_time_s":131.711,"test_pr_auc":0.6307183030345935,"_step":14,"test_loss":0.10761349875051364,"_wandb":{"runtime":1956},"_timestamp":1.7840929012398996e+09,"test_accuracy":0.9673010954962931,"train_vram_peak_mb":967.111328125,"val_f1":0.5570209464701319,"test_roc_auc":0.9193072775481443,"_runtime":1956,"global_step":63266,"pos_weight":10,"val_pr_auc":0.6136443695006331,"train_vram_allocated_mb":22.46337890625,"test_f1":0.5245374094931617,"val_roc_auc":0.9238255774138212,"val_loss":0.10250611246225581,"train_loss":0.1074708191769432,"val_accuracy":0.9684059093675649,"epoch":13}
\ No newline at end of file
diff --git a/wandb/run-20260714_234904-mj2wwse4/run-mj2wwse4.wandb b/wandb/run-20260714_234904-mj2wwse4/run-mj2wwse4.wandb
new file mode 100644
index 000000000..ee4958a47
Binary files /dev/null and b/wandb/run-20260714_234904-mj2wwse4/run-mj2wwse4.wandb differ
diff --git a/wandb/run-20260715_091251-mdnjqwik/files/config.yaml b/wandb/run-20260715_091251-mdnjqwik/files/config.yaml
new file mode 100644
index 000000000..a3a403b14
--- /dev/null
+++ b/wandb/run-20260715_091251-mdnjqwik/files/config.yaml
@@ -0,0 +1,228 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            1q8c2vgnjxiob0aki7xrpslgolctxedk:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --note-root
+                    - /shared/rsaas/physionet.org/files/mimic-note
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15979868160"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-15T14:12:51.592196Z"
+                writerId: 1q8c2vgnjxiob0aki7xrpslgolctxedk
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 16
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: /shared/rsaas/physionet.org/files/mimic-note
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: null
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260715_091251-mdnjqwik/files/requirements.txt b/wandb/run-20260715_091251-mdnjqwik/files/requirements.txt
new file mode 100644
index 000000000..d84484f0e
--- /dev/null
+++ b/wandb/run-20260715_091251-mdnjqwik/files/requirements.txt
@@ -0,0 +1,224 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
diff --git a/wandb/run-20260715_091251-mdnjqwik/files/wandb-metadata.json b/wandb/run-20260715_091251-mdnjqwik/files/wandb-metadata.json
new file mode 100644
index 000000000..faf84ac8a
--- /dev/null
+++ b/wandb/run-20260715_091251-mdnjqwik/files/wandb-metadata.json
@@ -0,0 +1,130 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-15T14:12:51.592196Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15979868160"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "1q8c2vgnjxiob0aki7xrpslgolctxedk"
+}
\ No newline at end of file
diff --git a/wandb/run-20260715_091251-mdnjqwik/files/wandb-summary.json b/wandb/run-20260715_091251-mdnjqwik/files/wandb-summary.json
new file mode 100644
index 000000000..2320ce951
--- /dev/null
+++ b/wandb/run-20260715_091251-mdnjqwik/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_wandb":{"runtime":8},"_runtime":8}
\ No newline at end of file
diff --git a/wandb/run-20260715_091251-mdnjqwik/run-mdnjqwik.wandb b/wandb/run-20260715_091251-mdnjqwik/run-mdnjqwik.wandb
new file mode 100644
index 000000000..b0993c60a
Binary files /dev/null and b/wandb/run-20260715_091251-mdnjqwik/run-mdnjqwik.wandb differ
diff --git a/wandb/run-20260715_100021-b532gynt/files/config.yaml b/wandb/run-20260715_100021-b532gynt/files/config.yaml
new file mode 100644
index 000000000..71dc27c4d
--- /dev/null
+++ b/wandb/run-20260715_100021-b532gynt/files/config.yaml
@@ -0,0 +1,228 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            u65w6ylqiu9k9sr922tbd8n6pj9y9pdn:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --note-root
+                    - /shared/rsaas/physionet.org/files/mimic-note
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 64
+                cpu_count_logical: 128
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "15981236224"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX A6000
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-7209ddee-de68-a17a-f65b-ac0f42256c43
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-b4e56019-e74d-298f-6c79-cdd64e15066e
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc
+                    - architecture: Ampere
+                      cudaCores: 10752
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX A6000
+                      uuid: GPU-12357f2c-fb85-9d97-90b0-a241510c022d
+                host: sunlab-serv-03.cs.illinois.edu
+                memory:
+                    total: "1081502597120"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.13
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-15T15:00:21.437575Z"
+                writerId: u65w6ylqiu9k9sr922tbd8n6pj9y9pdn
+        m: []
+        python_version: 3.12.13
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 16
+            "4": 3.12.13
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: /shared/rsaas/physionet.org/files/mimic-note
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: null
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260715_100021-b532gynt/files/requirements.txt b/wandb/run-20260715_100021-b532gynt/files/requirements.txt
new file mode 100644
index 000000000..d84484f0e
--- /dev/null
+++ b/wandb/run-20260715_100021-b532gynt/files/requirements.txt
@@ -0,0 +1,224 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
diff --git a/wandb/run-20260715_100021-b532gynt/files/wandb-metadata.json b/wandb/run-20260715_100021-b532gynt/files/wandb-metadata.json
new file mode 100644
index 000000000..734a33244
--- /dev/null
+++ b/wandb/run-20260715_100021-b532gynt/files/wandb-metadata.json
@@ -0,0 +1,130 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.13",
+  "startedAt":  "2026-07-15T15:00:21.437575Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "e24eec4a62ad0ad1ca875ad5fd100f36e1879e9b"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-serv-03.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  64,
+  "cpu_count_logical":  128,
+  "gpu":  "NVIDIA RTX A6000",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "15981236224"
+    }
+  },
+  "memory":  {
+    "total":  "1081502597120"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-29a27c7b-cd4b-9728-9cdc-7102f77d4548"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a34186de-ecc3-56c6-0e8f-ff8cfa0cc7b2"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-a733f3f6-b964-ca65-58ce-c7f32a13d7dc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-7209ddee-de68-a17a-f65b-ac0f42256c43"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fcf1f5ff-fc35-4cb4-fab0-dc9a5b6b6951"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-b4e56019-e74d-298f-6c79-cdd64e15066e"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-fbc6d802-bed5-9722-0929-a9cfd4d6b7bc"
+    },
+    {
+      "name":  "NVIDIA RTX A6000",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  10752,
+      "architecture":  "Ampere",
+      "uuid":  "GPU-12357f2c-fb85-9d97-90b0-a241510c022d"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "u65w6ylqiu9k9sr922tbd8n6pj9y9pdn"
+}
\ No newline at end of file
diff --git a/wandb/run-20260715_100021-b532gynt/files/wandb-summary.json b/wandb/run-20260715_100021-b532gynt/files/wandb-summary.json
new file mode 100644
index 000000000..2320ce951
--- /dev/null
+++ b/wandb/run-20260715_100021-b532gynt/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_wandb":{"runtime":8},"_runtime":8}
\ No newline at end of file
diff --git a/wandb/run-20260715_100021-b532gynt/run-b532gynt.wandb b/wandb/run-20260715_100021-b532gynt/run-b532gynt.wandb
new file mode 100644
index 000000000..f46f4d96c
Binary files /dev/null and b/wandb/run-20260715_100021-b532gynt/run-b532gynt.wandb differ
diff --git a/wandb/run-20260727_160331-uvc6o5zs/files/config.yaml b/wandb/run-20260727_160331-uvc6o5zs/files/config.yaml
new file mode 100644
index 000000000..a97cf5483
--- /dev/null
+++ b/wandb/run-20260727_160331-uvc6o5zs/files/config.yaml
@@ -0,0 +1,228 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            9k1337aubrwyt6p3hccyqsldgaosggxk:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /shared/eng/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "1e-3"
+                    - --weight-decay
+                    - "1e-5"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-only
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 128
+                cpu_count_logical: 255
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "9631342592"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 311a149fa5b0b01a248ce08d3c4ec66793cb9b6e
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX 6000 Ada Generation
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-25605dab-2bad-abf9-2f36-f0be6e776096
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-39ac7920-1d83-d4cf-c056-920c82599e52
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-6a135207-d247-4118-1218-1e2321919d87
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-583da23f-53e9-cf5f-2f63-7115778d21bf
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8
+                host: sunlab-c02.cs.illinois.edu
+                memory:
+                    total: "1081448284160"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-27T21:03:31.843087Z"
+                writerId: 9k1337aubrwyt6p3hccyqsldgaosggxk
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /shared/eng/wp14/pyhealth_cache_labs
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: null
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: null
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-only
+wandb_run_name:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260727_160331-uvc6o5zs/files/requirements.txt b/wandb/run-20260727_160331-uvc6o5zs/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260727_160331-uvc6o5zs/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260727_160331-uvc6o5zs/files/wandb-metadata.json b/wandb/run-20260727_160331-uvc6o5zs/files/wandb-metadata.json
new file mode 100644
index 000000000..5665d06a5
--- /dev/null
+++ b/wandb/run-20260727_160331-uvc6o5zs/files/wandb-metadata.json
@@ -0,0 +1,128 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-27T21:03:31.843087Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/shared/eng/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-only"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9631342592"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "9k1337aubrwyt6p3hccyqsldgaosggxk"
+}
\ No newline at end of file
diff --git a/wandb/run-20260727_160331-uvc6o5zs/files/wandb-summary.json b/wandb/run-20260727_160331-uvc6o5zs/files/wandb-summary.json
new file mode 100644
index 000000000..58f27969c
--- /dev/null
+++ b/wandb/run-20260727_160331-uvc6o5zs/files/wandb-summary.json
@@ -0,0 +1 @@
+{"test_pr_auc":0.6361709564058009,"val_roc_auc":0.9232587939709063,"epoch_time_s":147.761,"_runtime":3215,"pos_weight":10,"_timestamp":1.7851894287001421e+09,"test_accuracy":0.9668031426358304,"train_vram_allocated_mb":23.54931640625,"train_loss":0.1062656084620061,"val_pr_auc":0.6192477881133135,"epoch":19,"test_loss":0.10728481448527459,"test_roc_auc":0.9207827507695016,"test_f1":0.5238095238095238,"train_vram_peak_mb":968.197265625,"_wandb":{"runtime":3215},"global_step":90380,"val_loss":0.10286682227675893,"_step":20,"val_f1":0.5572998430141287,"val_accuracy":0.9687932274663863}
\ No newline at end of file
diff --git a/wandb/run-20260727_160331-uvc6o5zs/run-uvc6o5zs.wandb b/wandb/run-20260727_160331-uvc6o5zs/run-uvc6o5zs.wandb
new file mode 100644
index 000000000..18ade5b11
Binary files /dev/null and b/wandb/run-20260727_160331-uvc6o5zs/run-uvc6o5zs.wandb differ
diff --git a/wandb/run-20260727_181747-bz7db0ct/files/config.yaml b/wandb/run-20260727_181747-bz7db0ct/files/config.yaml
new file mode 100644
index 000000000..ed0406036
--- /dev/null
+++ b/wandb/run-20260727_181747-bz7db0ct/files/config.yaml
@@ -0,0 +1,235 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            ng73agh8nenn2b9jmunr6xvc35z4p1o4:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "64"
+                    - --hidden-dim
+                    - "64"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "1e-3"
+                    - --weight-decay
+                    - "1e-5"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-only
+                    - --wandb-run-name
+                    - labs_rnn_seed12
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 128
+                cpu_count_logical: 255
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "9632202752"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 311a149fa5b0b01a248ce08d3c4ec66793cb9b6e
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX 6000 Ada Generation
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-25605dab-2bad-abf9-2f36-f0be6e776096
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-39ac7920-1d83-d4cf-c056-920c82599e52
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-6a135207-d247-4118-1218-1e2321919d87
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-583da23f-53e9-cf5f-2f63-7115778d21bf
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8
+                host: sunlab-c02.cs.illinois.edu
+                memory:
+                    total: "1081448284160"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-27T23:17:47.555043Z"
+                writerId: ng73agh8nenn2b9jmunr6xvc35z4p1o4
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 64
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: null
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-only
+wandb_run_name:
+    value: labs_rnn_seed12
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260727_181747-bz7db0ct/files/requirements.txt b/wandb/run-20260727_181747-bz7db0ct/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260727_181747-bz7db0ct/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260727_181747-bz7db0ct/files/wandb-metadata.json b/wandb/run-20260727_181747-bz7db0ct/files/wandb-metadata.json
new file mode 100644
index 000000000..92a66fb1b
--- /dev/null
+++ b/wandb/run-20260727_181747-bz7db0ct/files/wandb-metadata.json
@@ -0,0 +1,132 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-27T23:17:47.555043Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "64",
+    "--hidden-dim",
+    "64",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "1",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-only",
+    "--wandb-run-name",
+    "labs_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9632202752"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "ng73agh8nenn2b9jmunr6xvc35z4p1o4"
+}
\ No newline at end of file
diff --git a/wandb/run-20260727_181747-bz7db0ct/files/wandb-summary.json b/wandb/run-20260727_181747-bz7db0ct/files/wandb-summary.json
new file mode 100644
index 000000000..bdf8c470f
--- /dev/null
+++ b/wandb/run-20260727_181747-bz7db0ct/files/wandb-summary.json
@@ -0,0 +1 @@
+{"epoch_time_s":125.353,"pos_weight":1,"_step":19,"test_f1":0.5792811839323467,"_wandb":{"runtime":2558},"test_pr_auc":0.6359263002799863,"val_f1":0.5513428120063191,"train_loss":0.10677212555209575,"val_loss":0.10355841470799879,"test_loss":0.10590764887103465,"val_roc_auc":0.9180975573391645,"val_pr_auc":0.6123994325554514,"train_vram_peak_mb":338.56591796875,"train_vram_allocated_mb":16.5615234375,"_timestamp":1.7851968266284647e+09,"global_step":85861,"epoch":18,"test_accuracy":0.9669691269226514,"val_accuracy":0.9685719028384884,"test_roc_auc":0.9212930525590202,"_runtime":2558}
\ No newline at end of file
diff --git a/wandb/run-20260727_181747-bz7db0ct/run-bz7db0ct.wandb b/wandb/run-20260727_181747-bz7db0ct/run-bz7db0ct.wandb
new file mode 100644
index 000000000..beaa316df
Binary files /dev/null and b/wandb/run-20260727_181747-bz7db0ct/run-bz7db0ct.wandb differ
diff --git a/wandb/run-20260728_084144-osiqkb7a/files/config.yaml b/wandb/run-20260728_084144-osiqkb7a/files/config.yaml
new file mode 100644
index 000000000..31d3b4c0d
--- /dev/null
+++ b/wandb/run-20260728_084144-osiqkb7a/files/config.yaml
@@ -0,0 +1,233 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            rgls97wpxhefl9injir90n0dc4vu2mzr:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs
+                    - --task
+                    - labs
+                    - --model
+                    - transformer
+                    - --embedding-dim
+                    - "64"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "1e-3"
+                    - --weight-decay
+                    - "1e-5"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --pos-weight
+                    - "1.0"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-only
+                    - --wandb-run-name
+                    - labs_transformer_seed12
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 128
+                cpu_count_logical: 255
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "9641213952"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 311a149fa5b0b01a248ce08d3c4ec66793cb9b6e
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX 6000 Ada Generation
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-25605dab-2bad-abf9-2f36-f0be6e776096
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-39ac7920-1d83-d4cf-c056-920c82599e52
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-6a135207-d247-4118-1218-1e2321919d87
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-583da23f-53e9-cf5f-2f63-7115778d21bf
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8
+                host: sunlab-c02.cs.illinois.edu
+                memory:
+                    total: "1081448284160"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-28T13:41:44.626252Z"
+                writerId: rgls97wpxhefl9injir90n0dc4vu2mzr
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 64
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: transformer
+note_root:
+    value: null
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-only
+wandb_run_name:
+    value: labs_transformer_seed12
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260728_084144-osiqkb7a/files/requirements.txt b/wandb/run-20260728_084144-osiqkb7a/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260728_084144-osiqkb7a/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260728_084144-osiqkb7a/files/wandb-metadata.json b/wandb/run-20260728_084144-osiqkb7a/files/wandb-metadata.json
new file mode 100644
index 000000000..931c0318c
--- /dev/null
+++ b/wandb/run-20260728_084144-osiqkb7a/files/wandb-metadata.json
@@ -0,0 +1,130 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-28T13:41:44.626252Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs",
+    "--task",
+    "labs",
+    "--model",
+    "transformer",
+    "--embedding-dim",
+    "64",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--pos-weight",
+    "1.0",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-only",
+    "--wandb-run-name",
+    "labs_transformer_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9641213952"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "rgls97wpxhefl9injir90n0dc4vu2mzr"
+}
\ No newline at end of file
diff --git a/wandb/run-20260728_084144-osiqkb7a/files/wandb-summary.json b/wandb/run-20260728_084144-osiqkb7a/files/wandb-summary.json
new file mode 100644
index 000000000..dd38351d5
--- /dev/null
+++ b/wandb/run-20260728_084144-osiqkb7a/files/wandb-summary.json
@@ -0,0 +1 @@
+{"train_vram_peak_mb":27835.14208984375,"_runtime":1938,"val_pr_auc":0.5687590485828822,"val_roc_auc":0.9159081391282424,"test_pr_auc":0.6129330211587669,"epoch":15,"val_loss":0.11564977598763937,"test_f1":0.5586510263929618,"test_roc_auc":0.9144428282915618,"pos_weight":1,"global_step":72304,"train_loss":0.11545573439666292,"test_loss":0.11046364758394461,"val_accuracy":0.9658053449897638,"train_vram_allocated_mb":17.42041015625,"val_f1":0.5179407176287052,"test_accuracy":0.9666924864446166,"_step":16,"epoch_time_s":111.854,"_timestamp":1.7852480443924367e+09,"_wandb":{"runtime":1938}}
\ No newline at end of file
diff --git a/wandb/run-20260728_084144-osiqkb7a/run-osiqkb7a.wandb b/wandb/run-20260728_084144-osiqkb7a/run-osiqkb7a.wandb
new file mode 100644
index 000000000..34b0071d4
Binary files /dev/null and b/wandb/run-20260728_084144-osiqkb7a/run-osiqkb7a.wandb differ
diff --git a/wandb/run-20260728_123606-i66z42nf/files/config.yaml b/wandb/run-20260728_123606-i66z42nf/files/config.yaml
new file mode 100644
index 000000000..6b64c9ed8
--- /dev/null
+++ b/wandb/run-20260728_123606-i66z42nf/files/config.yaml
@@ -0,0 +1,233 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            cnrv3ntwx5zyzo4cvx2pxsk1p03yy39d:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --note-root
+                    - /shared/rsaas/physionet.org/files/mimic-note
+                    - --cache-dir
+                    - /home/wp14/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "1e-3"
+                    - --weight-decay
+                    - "1e-5"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --wandb-run-name
+                    - labs_notes_rnn_seed12
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 128
+                cpu_count_logical: 255
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "9642340352"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 311a149fa5b0b01a248ce08d3c4ec66793cb9b6e
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX 6000 Ada Generation
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-25605dab-2bad-abf9-2f36-f0be6e776096
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-39ac7920-1d83-d4cf-c056-920c82599e52
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-6a135207-d247-4118-1218-1e2321919d87
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-583da23f-53e9-cf5f-2f63-7115778d21bf
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8
+                host: sunlab-c02.cs.illinois.edu
+                memory:
+                    total: "1081448284160"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-28T17:36:06.511561Z"
+                writerId: cnrv3ntwx5zyzo4cvx2pxsk1p03yy39d
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 15
+                - 16
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/wp14/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: false
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: /shared/rsaas/physionet.org/files/mimic-note
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: null
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: labs_notes_rnn_seed12
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260728_123606-i66z42nf/files/requirements.txt b/wandb/run-20260728_123606-i66z42nf/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260728_123606-i66z42nf/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260728_123606-i66z42nf/files/wandb-metadata.json b/wandb/run-20260728_123606-i66z42nf/files/wandb-metadata.json
new file mode 100644
index 000000000..07613567a
--- /dev/null
+++ b/wandb/run-20260728_123606-i66z42nf/files/wandb-metadata.json
@@ -0,0 +1,132 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-28T17:36:06.511561Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/home/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--wandb-run-name",
+    "labs_notes_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9642340352"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "cnrv3ntwx5zyzo4cvx2pxsk1p03yy39d"
+}
\ No newline at end of file
diff --git a/wandb/run-20260728_123606-i66z42nf/files/wandb-summary.json b/wandb/run-20260728_123606-i66z42nf/files/wandb-summary.json
new file mode 100644
index 000000000..f60bccc79
--- /dev/null
+++ b/wandb/run-20260728_123606-i66z42nf/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_wandb":{"runtime":6},"_runtime":6}
\ No newline at end of file
diff --git a/wandb/run-20260728_123606-i66z42nf/run-i66z42nf.wandb b/wandb/run-20260728_123606-i66z42nf/run-i66z42nf.wandb
new file mode 100644
index 000000000..aaef3d7f9
Binary files /dev/null and b/wandb/run-20260728_123606-i66z42nf/run-i66z42nf.wandb differ
diff --git a/wandb/run-20260728_175103-bd9dek9x/files/requirements.txt b/wandb/run-20260728_175103-bd9dek9x/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260728_175103-bd9dek9x/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260728_175103-bd9dek9x/files/wandb-metadata.json b/wandb/run-20260728_175103-bd9dek9x/files/wandb-metadata.json
new file mode 100644
index 000000000..edfa409ae
--- /dev/null
+++ b/wandb/run-20260728_175103-bd9dek9x/files/wandb-metadata.json
@@ -0,0 +1,133 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-28T22:51:03.133326Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/shared/rsaas/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--wandb-run-name",
+    "labs_notes_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9620135936"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "hywo7zaph59r2wyxsjmse4w92d23z449"
+}
\ No newline at end of file
diff --git a/wandb/run-20260728_175103-bd9dek9x/run-bd9dek9x.wandb b/wandb/run-20260728_175103-bd9dek9x/run-bd9dek9x.wandb
new file mode 100644
index 000000000..013c097ee
Binary files /dev/null and b/wandb/run-20260728_175103-bd9dek9x/run-bd9dek9x.wandb differ
diff --git a/wandb/run-20260729_103813-pyxpnamr/files/requirements.txt b/wandb/run-20260729_103813-pyxpnamr/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260729_103813-pyxpnamr/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260729_103813-pyxpnamr/files/wandb-metadata.json b/wandb/run-20260729_103813-pyxpnamr/files/wandb-metadata.json
new file mode 100644
index 000000000..781464d36
--- /dev/null
+++ b/wandb/run-20260729_103813-pyxpnamr/files/wandb-metadata.json
@@ -0,0 +1,135 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-29T15:38:13.490309Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/shared/rsaas/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--pos-weight",
+    "1",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--wandb-run-name",
+    "labs_notes_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9632387072"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "721tzyup2w7a8zl4arii4sj0atrkjgjz"
+}
\ No newline at end of file
diff --git a/wandb/run-20260729_103813-pyxpnamr/run-pyxpnamr.wandb b/wandb/run-20260729_103813-pyxpnamr/run-pyxpnamr.wandb
new file mode 100644
index 000000000..b6e1a363b
Binary files /dev/null and b/wandb/run-20260729_103813-pyxpnamr/run-pyxpnamr.wandb differ
diff --git a/wandb/run-20260729_230052-vhbvau7e/files/config.yaml b/wandb/run-20260729_230052-vhbvau7e/files/config.yaml
new file mode 100644
index 000000000..2694a2f30
--- /dev/null
+++ b/wandb/run-20260729_230052-vhbvau7e/files/config.yaml
@@ -0,0 +1,238 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            kify0s9esdi6ixsxxc3mz9o03nqi3qa8:
+                args:
+                    - --ehr-root
+                    - /shared/rsaas/physionet.org/files/mimiciv/2.2
+                    - --note-root
+                    - /shared/rsaas/physionet.org/files/mimic-note
+                    - --cache-dir
+                    - /shared/rsaas/wp14/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --epochs
+                    - "15"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "1e-3"
+                    - --weight-decay
+                    - "1e-5"
+                    - --patience
+                    - "5"
+                    - --num-workers
+                    - "4"
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/wp14/output
+                    - --pos-weight
+                    - "1"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --wandb-run-name
+                    - labs_notes_rnn_seed12
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 128
+                cpu_count_logical: 255
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "34342961152"
+                        used: "9641287680"
+                email: williampangbest1@gmail.com
+                executable: /home/wp14/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 311a149fa5b0b01a248ce08d3c4ec66793cb9b6e
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA RTX 6000 Ada Generation
+                gpu_count: 8
+                gpu_nvidia:
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-25605dab-2bad-abf9-2f36-f0be6e776096
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-39ac7920-1d83-d4cf-c056-920c82599e52
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-6a135207-d247-4118-1218-1e2321919d87
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-583da23f-53e9-cf5f-2f63-7115778d21bf
+                    - architecture: Ada
+                      cudaCores: 18176
+                      memoryTotal: "51527024640"
+                      name: NVIDIA RTX 6000 Ada Generation
+                      uuid: GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8
+                host: sunlab-c02.cs.illinois.edu
+                memory:
+                    total: "1081448284160"
+                os: Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28
+                program: /home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/wp14/PyHealth
+                startedAt: "2026-07-30T04:00:52.025825Z"
+                writerId: kify0s9esdi6ixsxxc3mz9o03nqi3qa8
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /shared/rsaas/wp14/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /shared/rsaas/physionet.org/files/mimiciv/2.2
+embedding_dim:
+    value: 128
+epochs:
+    value: 15
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: /shared/rsaas/physionet.org/files/mimic-note
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/wp14/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: labs_notes_rnn_seed12
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260729_230052-vhbvau7e/files/requirements.txt b/wandb/run-20260729_230052-vhbvau7e/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260729_230052-vhbvau7e/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260729_230052-vhbvau7e/files/wandb-metadata.json b/wandb/run-20260729_230052-vhbvau7e/files/wandb-metadata.json
new file mode 100644
index 000000000..905fcd20e
--- /dev/null
+++ b/wandb/run-20260729_230052-vhbvau7e/files/wandb-metadata.json
@@ -0,0 +1,135 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-07-30T04:00:52.025825Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/shared/rsaas/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "15",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--pos-weight",
+    "1",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--wandb-run-name",
+    "labs_notes_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "311a149fa5b0b01a248ce08d3c4ec66793cb9b6e"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9641287680"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "kify0s9esdi6ixsxxc3mz9o03nqi3qa8"
+}
\ No newline at end of file
diff --git a/wandb/run-20260729_230052-vhbvau7e/files/wandb-summary.json b/wandb/run-20260729_230052-vhbvau7e/files/wandb-summary.json
new file mode 100644
index 000000000..001066980
--- /dev/null
+++ b/wandb/run-20260729_230052-vhbvau7e/files/wandb-summary.json
@@ -0,0 +1 @@
+{"val_f1":0.072480181200453,"_wandb":{"runtime":300713},"train_loss":0.16350290598687833,"epoch":13,"val_pr_auc":0.22141367898844128,"test_accuracy":0.9521965253955958,"test_f1":0.048458149779735685,"val_loss":0.1570026268506736,"train_vram_allocated_mb":438.35595703125,"pos_weight":1,"global_step":63266,"val_roc_auc":0.7993363615862891,"_timestamp":1.7856847667407796e+09,"train_vram_peak_mb":12370.1982421875,"test_loss":0.16382228572547963,"epoch_time_s":6052.221,"test_roc_auc":0.791545657966757,"_step":14,"_runtime":300713,"val_accuracy":0.9546837824378908,"test_pr_auc":0.23978078766730027}
\ No newline at end of file
diff --git a/wandb/run-20260729_230052-vhbvau7e/run-vhbvau7e.wandb b/wandb/run-20260729_230052-vhbvau7e/run-vhbvau7e.wandb
new file mode 100644
index 000000000..746787a9c
Binary files /dev/null and b/wandb/run-20260729_230052-vhbvau7e/run-vhbvau7e.wandb differ
diff --git a/wandb/run-20260803_062203-t13ix6uv/files/config.yaml b/wandb/run-20260803_062203-t13ix6uv/files/config.yaml
new file mode 100644
index 000000000..5262314d1
--- /dev/null
+++ b/wandb/run-20260803_062203-t13ix6uv/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            ehnpm8n0mxxv1fhcayqyos8u60mb7fch:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - ehrmamba
+                    - --embedding-dim
+                    - "96"
+                    - --num-layers
+                    - "2"
+                    - --mamba-state-size
+                    - "16"
+                    - --mamba-conv-kernel
+                    - "4"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "2"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "39855661056"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 8ad2d8b82c5e8eb03bf020984a4669f9aa01d5f8
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153
+                host: 192-222-55-48
+                memory:
+                    total: "237490823168"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-03T06:22:03.015199Z"
+                writerId: ehnpm8n0mxxv1fhcayqyos8u60mb7fch
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 2
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 96
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: ehrmamba
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260803_062203-t13ix6uv/files/requirements.txt b/wandb/run-20260803_062203-t13ix6uv/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260803_062203-t13ix6uv/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260803_062203-t13ix6uv/files/wandb-metadata.json b/wandb/run-20260803_062203-t13ix6uv/files/wandb-metadata.json
new file mode 100644
index 000000000..09242eba4
--- /dev/null
+++ b/wandb/run-20260803_062203-t13ix6uv/files/wandb-metadata.json
@@ -0,0 +1,87 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-03T06:22:03.015199Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "ehrmamba",
+    "--embedding-dim",
+    "96",
+    "--num-layers",
+    "2",
+    "--mamba-state-size",
+    "16",
+    "--mamba-conv-kernel",
+    "4",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "2",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "8ad2d8b82c5e8eb03bf020984a4669f9aa01d5f8"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-55-48",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "39855661056"
+    }
+  },
+  "memory":  {
+    "total":  "237490823168"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "ehnpm8n0mxxv1fhcayqyos8u60mb7fch"
+}
\ No newline at end of file
diff --git a/wandb/run-20260803_062203-t13ix6uv/files/wandb-summary.json b/wandb/run-20260803_062203-t13ix6uv/files/wandb-summary.json
new file mode 100644
index 000000000..31959d39d
--- /dev/null
+++ b/wandb/run-20260803_062203-t13ix6uv/files/wandb-summary.json
@@ -0,0 +1 @@
+{"test_pr_auc":0.2257047616279328,"train_loss":0.2407169388210231,"test_roc_auc":0.8039630828452551,"_wandb":{"runtime":10511},"val_loss":0.19066320987775076,"train_vram_peak_mb":996.23291015625,"test_accuracy":0.9514772601527055,"val_f1":0,"epoch_time_s":1186.134,"_runtime":10511,"test_loss":0.25312376624519245,"val_accuracy":0.9542411331820948,"test_f1":0,"_timestamp":1.78574863531726e+09,"epoch":7,"global_step":578344,"val_roc_auc":0.813408566359113,"_step":8,"pos_weight":1,"train_vram_allocated_mb":483.44287109375,"val_pr_auc":0.22149906687571344}
\ No newline at end of file
diff --git a/wandb/run-20260803_062203-t13ix6uv/run-t13ix6uv.wandb b/wandb/run-20260803_062203-t13ix6uv/run-t13ix6uv.wandb
new file mode 100644
index 000000000..71cf312fd
Binary files /dev/null and b/wandb/run-20260803_062203-t13ix6uv/run-t13ix6uv.wandb differ
diff --git a/wandb/run-20260803_141932-fyilcb4j/files/config.yaml b/wandb/run-20260803_141932-fyilcb4j/files/config.yaml
new file mode 100644
index 000000000..b8d381396
--- /dev/null
+++ b/wandb/run-20260803_141932-fyilcb4j/files/config.yaml
@@ -0,0 +1,212 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            nbk8a0le70n2se7zjernlaez70luwqai:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - jambaehr
+                    - --embedding-dim
+                    - "96"
+                    - --heads
+                    - "4"
+                    - --jamba-transformer-layers
+                    - "2"
+                    - --jamba-mamba-layers
+                    - "6"
+                    - --mamba-state-size
+                    - "16"
+                    - --mamba-conv-kernel
+                    - "4"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "2"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "40472424448"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: f9d669fdacb65f4336adf784d674a90bd2fb13fe
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153
+                host: 192-222-55-48
+                memory:
+                    total: "237490823168"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-03T14:19:32.420352Z"
+                writerId: nbk8a0le70n2se7zjernlaez70luwqai
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 2
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 96
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: jambaehr
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260803_141932-fyilcb4j/files/requirements.txt b/wandb/run-20260803_141932-fyilcb4j/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260803_141932-fyilcb4j/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260803_141932-fyilcb4j/files/wandb-metadata.json b/wandb/run-20260803_141932-fyilcb4j/files/wandb-metadata.json
new file mode 100644
index 000000000..abe4e77b7
--- /dev/null
+++ b/wandb/run-20260803_141932-fyilcb4j/files/wandb-metadata.json
@@ -0,0 +1,91 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-03T14:19:32.420352Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "jambaehr",
+    "--embedding-dim",
+    "96",
+    "--heads",
+    "4",
+    "--jamba-transformer-layers",
+    "2",
+    "--jamba-mamba-layers",
+    "6",
+    "--mamba-state-size",
+    "16",
+    "--mamba-conv-kernel",
+    "4",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "2",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "f9d669fdacb65f4336adf784d674a90bd2fb13fe"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-55-48",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "40472424448"
+    }
+  },
+  "memory":  {
+    "total":  "237490823168"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "nbk8a0le70n2se7zjernlaez70luwqai"
+}
\ No newline at end of file
diff --git a/wandb/run-20260803_141932-fyilcb4j/files/wandb-summary.json b/wandb/run-20260803_141932-fyilcb4j/files/wandb-summary.json
new file mode 100644
index 000000000..a67b1e378
--- /dev/null
+++ b/wandb/run-20260803_141932-fyilcb4j/files/wandb-summary.json
@@ -0,0 +1 @@
+{"epoch":7,"_runtime":18377,"_timestamp":1.7857851505828013e+09,"test_loss":0.18805167121359045,"epoch_time_s":2146.508,"_wandb":{"runtime":18377},"val_loss":0.15788743355858376,"val_pr_auc":0.2928875906699203,"test_roc_auc":0.8442603119037296,"pos_weight":1,"global_step":578344,"val_roc_auc":0.8371692940101002,"val_f1":0.11173184357541899,"train_vram_peak_mb":1008.55810546875,"test_accuracy":0.9528051344472723,"test_f1":0.13925327951564076,"train_vram_allocated_mb":495.76806640625,"train_loss":0.1696291825593504,"test_pr_auc":0.2971645611465631,"_step":8,"val_accuracy":0.9560117302052786}
\ No newline at end of file
diff --git a/wandb/run-20260803_141932-fyilcb4j/run-fyilcb4j.wandb b/wandb/run-20260803_141932-fyilcb4j/run-fyilcb4j.wandb
new file mode 100644
index 000000000..9581e4080
Binary files /dev/null and b/wandb/run-20260803_141932-fyilcb4j/run-fyilcb4j.wandb differ
diff --git a/wandb/run-20260803_193000-2shg7x9f/files/config.yaml b/wandb/run-20260803_193000-2shg7x9f/files/config.yaml
new file mode 100644
index 000000000..821cedd1d
--- /dev/null
+++ b/wandb/run-20260803_193000-2shg7x9f/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            8tvn49su1dafvrzz7igf1e5liy9sa0kl:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - bottleneck_transformer
+                    - --embedding-dim
+                    - "128"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --bottlenecks-n
+                    - "4"
+                    - --fusion-startidx
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "40507023360"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: f9d669fdacb65f4336adf784d674a90bd2fb13fe
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153
+                host: 192-222-55-48
+                memory:
+                    total: "237490823168"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-03T19:30:00.341184Z"
+                writerId: 8tvn49su1dafvrzz7igf1e5liy9sa0kl
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 15
+                - 16
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: bottleneck_transformer
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260803_193000-2shg7x9f/files/requirements.txt b/wandb/run-20260803_193000-2shg7x9f/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260803_193000-2shg7x9f/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260803_193000-2shg7x9f/files/wandb-metadata.json b/wandb/run-20260803_193000-2shg7x9f/files/wandb-metadata.json
new file mode 100644
index 000000000..a94aca378
--- /dev/null
+++ b/wandb/run-20260803_193000-2shg7x9f/files/wandb-metadata.json
@@ -0,0 +1,89 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-03T19:30:00.341184Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "bottleneck_transformer",
+    "--embedding-dim",
+    "128",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--bottlenecks-n",
+    "4",
+    "--fusion-startidx",
+    "1",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "f9d669fdacb65f4336adf784d674a90bd2fb13fe"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-55-48",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "40507023360"
+    }
+  },
+  "memory":  {
+    "total":  "237490823168"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "8tvn49su1dafvrzz7igf1e5liy9sa0kl"
+}
\ No newline at end of file
diff --git a/wandb/run-20260803_193000-2shg7x9f/files/wandb-summary.json b/wandb/run-20260803_193000-2shg7x9f/files/wandb-summary.json
new file mode 100644
index 000000000..50a3b558f
--- /dev/null
+++ b/wandb/run-20260803_193000-2shg7x9f/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_wandb":{"runtime":1707},"_runtime":1707}
\ No newline at end of file
diff --git a/wandb/run-20260803_193000-2shg7x9f/run-2shg7x9f.wandb b/wandb/run-20260803_193000-2shg7x9f/run-2shg7x9f.wandb
new file mode 100644
index 000000000..6ab095f3d
Binary files /dev/null and b/wandb/run-20260803_193000-2shg7x9f/run-2shg7x9f.wandb differ
diff --git a/wandb/run-20260803_203631-2jxwxgpk/files/config.yaml b/wandb/run-20260803_203631-2jxwxgpk/files/config.yaml
new file mode 100644
index 000000000..4012c2d00
--- /dev/null
+++ b/wandb/run-20260803_203631-2jxwxgpk/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            lv9zun5gwcke312qqebfx0htwufrxqup:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - bottleneck_transformer
+                    - --embedding-dim
+                    - "128"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --bottlenecks-n
+                    - "4"
+                    - --fusion-startidx
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "40509689856"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: f9d669fdacb65f4336adf784d674a90bd2fb13fe
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153
+                host: 192-222-55-48
+                memory:
+                    total: "237490823168"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-03T20:36:31.314140Z"
+                writerId: lv9zun5gwcke312qqebfx0htwufrxqup
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 15
+                - 16
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: bottleneck_transformer
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260803_203631-2jxwxgpk/files/requirements.txt b/wandb/run-20260803_203631-2jxwxgpk/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260803_203631-2jxwxgpk/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260803_203631-2jxwxgpk/files/wandb-metadata.json b/wandb/run-20260803_203631-2jxwxgpk/files/wandb-metadata.json
new file mode 100644
index 000000000..1037a2377
--- /dev/null
+++ b/wandb/run-20260803_203631-2jxwxgpk/files/wandb-metadata.json
@@ -0,0 +1,89 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-03T20:36:31.314140Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "bottleneck_transformer",
+    "--embedding-dim",
+    "128",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--bottlenecks-n",
+    "4",
+    "--fusion-startidx",
+    "1",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "f9d669fdacb65f4336adf784d674a90bd2fb13fe"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-55-48",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "40509689856"
+    }
+  },
+  "memory":  {
+    "total":  "237490823168"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "lv9zun5gwcke312qqebfx0htwufrxqup"
+}
\ No newline at end of file
diff --git a/wandb/run-20260803_203631-2jxwxgpk/files/wandb-summary.json b/wandb/run-20260803_203631-2jxwxgpk/files/wandb-summary.json
new file mode 100644
index 000000000..a53e3c8e2
--- /dev/null
+++ b/wandb/run-20260803_203631-2jxwxgpk/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_runtime":1671,"_wandb":{"runtime":1671}}
\ No newline at end of file
diff --git a/wandb/run-20260803_203631-2jxwxgpk/run-2jxwxgpk.wandb b/wandb/run-20260803_203631-2jxwxgpk/run-2jxwxgpk.wandb
new file mode 100644
index 000000000..2985f09cf
Binary files /dev/null and b/wandb/run-20260803_203631-2jxwxgpk/run-2jxwxgpk.wandb differ
diff --git a/wandb/run-20260803_211730-5pkaxkjn/files/config.yaml b/wandb/run-20260803_211730-5pkaxkjn/files/config.yaml
new file mode 100644
index 000000000..7aac7b7b8
--- /dev/null
+++ b/wandb/run-20260803_211730-5pkaxkjn/files/config.yaml
@@ -0,0 +1,208 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            dakcux8zokfcev2ptfnwbp0xjacknbp9:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+                    - --task
+                    - notes_labs
+                    - --model
+                    - bottleneck_transformer
+                    - --embedding-dim
+                    - "128"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --bottlenecks-n
+                    - "4"
+                    - --fusion-startidx
+                    - "1"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "40512843776"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: f9d669fdacb65f4336adf784d674a90bd2fb13fe
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153
+                host: 192-222-55-48
+                memory:
+                    total: "237490823168"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-03T21:17:30.385541Z"
+                writerId: dakcux8zokfcev2ptfnwbp0xjacknbp9
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 13
+                - 15
+                - 16
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: bottleneck_transformer
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260803_211730-5pkaxkjn/files/requirements.txt b/wandb/run-20260803_211730-5pkaxkjn/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260803_211730-5pkaxkjn/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260803_211730-5pkaxkjn/files/wandb-metadata.json b/wandb/run-20260803_211730-5pkaxkjn/files/wandb-metadata.json
new file mode 100644
index 000000000..e4b85a0fb
--- /dev/null
+++ b/wandb/run-20260803_211730-5pkaxkjn/files/wandb-metadata.json
@@ -0,0 +1,89 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-03T21:17:30.385541Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "bottleneck_transformer",
+    "--embedding-dim",
+    "128",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--bottlenecks-n",
+    "4",
+    "--fusion-startidx",
+    "1",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "f9d669fdacb65f4336adf784d674a90bd2fb13fe"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-55-48",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "40512843776"
+    }
+  },
+  "memory":  {
+    "total":  "237490823168"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-f6e0d80a-1ee4-6af1-6585-7b23e7185153"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "dakcux8zokfcev2ptfnwbp0xjacknbp9"
+}
\ No newline at end of file
diff --git a/wandb/run-20260803_211730-5pkaxkjn/files/wandb-summary.json b/wandb/run-20260803_211730-5pkaxkjn/files/wandb-summary.json
new file mode 100644
index 000000000..ad2d283e6
--- /dev/null
+++ b/wandb/run-20260803_211730-5pkaxkjn/files/wandb-summary.json
@@ -0,0 +1 @@
+{"_runtime":1674,"_wandb":{"runtime":1674}}
\ No newline at end of file
diff --git a/wandb/run-20260803_211730-5pkaxkjn/run-5pkaxkjn.wandb b/wandb/run-20260803_211730-5pkaxkjn/run-5pkaxkjn.wandb
new file mode 100644
index 000000000..4c4bf7a5e
Binary files /dev/null and b/wandb/run-20260803_211730-5pkaxkjn/run-5pkaxkjn.wandb differ
diff --git a/wandb/run-20260806_100316-vs4k3zp2/files/requirements.txt b/wandb/run-20260806_100316-vs4k3zp2/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260806_100316-vs4k3zp2/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260806_100316-vs4k3zp2/files/wandb-metadata.json b/wandb/run-20260806_100316-vs4k3zp2/files/wandb-metadata.json
new file mode 100644
index 000000000..4d036037c
--- /dev/null
+++ b/wandb/run-20260806_100316-vs4k3zp2/files/wandb-metadata.json
@@ -0,0 +1,135 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-06T15:03:16.148912Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cache-dir",
+    "/shared/rsaas/wp14/pyhealth_cache_labs_notes",
+    "--task",
+    "notes_labs",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "15",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--pos-weight",
+    "1",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes",
+    "--wandb-run-name",
+    "labs_notes_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "a83f7d6a1dbbc9ff710a88d6027a332e84fde041"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9714499584"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "nurt0tm2lyca5y2qweo4zg5doaucagzs"
+}
\ No newline at end of file
diff --git a/wandb/run-20260806_100316-vs4k3zp2/run-vs4k3zp2.wandb b/wandb/run-20260806_100316-vs4k3zp2/run-vs4k3zp2.wandb
new file mode 100644
index 000000000..30977402b
Binary files /dev/null and b/wandb/run-20260806_100316-vs4k3zp2/run-vs4k3zp2.wandb differ
diff --git a/wandb/run-20260806_113219-dsohsz9y/files/requirements.txt b/wandb/run-20260806_113219-dsohsz9y/files/requirements.txt
new file mode 100644
index 000000000..f73b3efe9
--- /dev/null
+++ b/wandb/run-20260806_113219-dsohsz9y/files/requirements.txt
@@ -0,0 +1,226 @@
+packaging==26.0
+setuptools==82.0.1
+wheel==0.46.3
+pip==26.0.1
+tblib==3.2.2
+sympy==1.14.0
+mpmath==1.3.0
+typing_extensions==4.15.0
+six==1.17.0
+pillow==12.1.1
+nvidia-nvjitlink-cu12==12.6.85
+nvidia-curand-cu12==10.3.7.77
+numpy==2.2.6
+nvidia-cufft-cu12==11.3.0.4
+nvidia-cuda-runtime-cu12==12.6.77
+nvidia-cuda-nvrtc-cu12==12.6.77
+nvidia-cuda-cupti-cu12==12.6.80
+nvidia-cublas-cu12==12.6.4.1
+narwhals==2.13.0
+networkx==3.6.1
+MarkupSafe==3.0.3
+fsspec==2026.2.0
+filelock==3.25.2
+nvidia-cusparse-cu12==12.5.4.2
+nvidia-cudnn-cu12==9.5.1.17
+Jinja2==3.1.6
+nvidia-cusolver-cu12==11.7.1.2
+torch-einops-utils==0.0.30
+sortedcontainers==2.4.0
+pytz==2026.1.post1
+nvidia-cusparselt-cu12==0.6.3
+zict==3.0.0
+xyzservices==2026.3.0
+urllib3==2.5.0
+urllib3==2.7.0
+tzdata==2026.1
+tzdata==2026.2
+typing-inspection==0.4.2
+safetensors==0.7.0
+triton==3.3.1
+tqdm==4.67.3
+tornado==6.5.5
+tornado==6.5.7
+toolz==1.1.0
+threadpoolctl==3.6.0
+regex==2026.4.4
+PyYAML==6.0.3
+pyparsing==3.3.2
+pydantic_core==2.33.2
+pydantic_core==2.46.4
+pyarrow==22.0.0
+psutil==7.2.2
+polars-runtime-32==1.35.2
+platformdirs==4.9.6
+platformdirs==4.10.0
+obstore==0.9.2
+nvidia-nvtx-cu12==12.6.77
+nvidia-nccl-cu12==2.26.2
+nvidia-cufile-cu12==1.11.1.6
+msgpack==1.1.2
+more-itertools==10.8.0
+lz4==4.4.5
+locket==1.0.0
+littleutils==0.2.4
+lightning-utilities==0.15.3
+lazy-loader==0.5
+kiwisolver==1.5.0
+joblib==1.5.3
+jmespath==1.1.0
+idna==3.11
+idna==3.18
+hf-xet==1.4.3
+fonttools==4.62.1
+einops==0.8.2
+decorator==5.2.1
+decorator==5.3.1
+cycler==0.12.1
+cloudpickle==3.1.2
+click==8.3.2
+click==8.4.2
+charset-normalizer==3.4.7
+certifi==2026.2.25
+certifi==2026.6.17
+annotated-types==0.7.0
+tifffile==2026.3.3
+scipy==1.17.1
+requests==2.33.1
+requests==2.34.2
+rdkit==2026.3.1
+python-dateutil==2.9.0.post0
+pydantic==2.11.10
+pydantic==2.13.4
+polars==1.35.2
+partd==1.4.2
+contourpy==1.3.3
+scikit-learn==1.7.2
+pooch==1.9.0
+pandas==2.3.3
+outdated==0.2.2
+matplotlib==3.10.8
+huggingface_hub==0.36.2
+dask==2025.11.0
+botocore==1.42.88
+bokeh==3.9.0
+torch==2.7.1
+tokenizers==0.21.4
+s3transfer==0.16.0
+mne==1.10.2
+distributed==2025.11.0
+transformers==4.53.3
+pyhealth==2.0.0
+torchvision==0.22.1
+ogb==1.3.6
+linformer==0.2.3
+boto3==1.42.88
+axial_positional_embedding==0.3.12
+accelerate==1.13.0
+peft==0.18.1
+litdata==0.2.61
+hyper-connections==0.4.9
+local-attention==1.11.2
+CoLT5-attention==0.11.1
+product_key_memory==0.3.0
+linear-attention-transformer==0.19.1
+nvidia-ml-py==13.590.48
+nvitop==1.7.0
+asttokens==3.0.1
+attrs==26.1.0
+babel==2.18.0
+backports.zstd==1.6.0
+Brotli==1.2.0
+cached-property==1.5.2
+comm==0.2.3
+debugpy==1.8.21
+defusedxml==0.7.1
+executing==2.2.1
+hpack==4.1.0
+hyperframe==6.1.0
+json5==0.15.0
+jsonpointer==3.1.1
+lark==1.3.1
+nest-asyncio2==1.7.2
+pandocfilters==1.5.0
+parso==0.8.7
+prometheus_client==0.25.0
+ptyprocess==0.7.0
+pure_eval==0.2.3
+pycparser==3.0
+Pygments==2.20.0
+PySocks==1.7.1
+fastjsonschema==2.21.2
+rfc3986-validator==0.1.1
+rpds-py==2026.5.1
+Send2Trash==2.1.0
+sniffio==1.3.1
+soupsieve==2.8.4
+tomli==2.4.1
+traitlets==5.15.1
+typing_utils==0.1.0
+uri-template==1.3.0
+wcwidth==0.8.1
+webcolors==25.10.0
+webencodings==0.5.1
+websocket-client==1.9.0
+zipp==4.1.0
+async-lru==2.3.0
+bleach==6.4.0
+cffi==1.17.1
+exceptiongroup==1.3.1
+h11==0.16.0
+h2==4.3.0
+importlib_metadata==9.0.0
+ipython_pygments_lexers==1.1.1
+jedi==0.19.2
+jupyter_core==5.9.1
+jupyterlab_pygments==0.3.0
+matplotlib-inline==0.2.2
+mistune==3.3.2
+overrides==7.7.0
+pexpect==4.9.0
+prompt_toolkit==3.0.52
+python-json-logger==4.1.0
+referencing==0.37.0
+rfc3339_validator==0.1.4
+rfc3987-syntax==1.1.0
+stack_data==0.6.3
+terminado==0.18.1
+tinycss2==1.4.0
+anyio==4.14.0
+argon2-cffi-bindings==25.1.0
+arrow==1.4.0
+beautifulsoup4==4.15.0
+fqdn==1.5.1
+ipython==9.14.1
+jsonschema-specifications==2025.9.1
+jupyter_builder==1.0.2
+jupyter_server_terminals==0.5.4
+argon2-cffi==25.1.0
+httpcore==1.0.9
+isoduration==20.11.0
+jsonschema==4.26.0
+pyzmq==27.1.0
+httpx==0.28.1
+jupyter_client==8.9.1
+nbformat==5.10.4
+ipykernel==7.3.0
+jupyter-events==0.12.1
+nbclient==0.11.0
+nbconvert==7.17.1
+jupyter_server==2.20.0
+jupyter-lsp==2.3.1
+jupyterlab_server==2.28.0
+notebook_shim==0.2.4
+jupyterlab==4.6.0
+notebook==7.6.0
+protobuf==6.33.5
+smmap==5.0.3
+gitdb==4.0.12
+GitPython==3.1.50
+sentry-sdk==2.64.0
+wandb==0.28.0
+widgetsnbextension==4.0.15
+jupyterlab_widgets==3.0.16
+ipywidgets==8.1.8
+htcondor==25.11.0
+htcondor-cli==25.11.0
diff --git a/wandb/run-20260806_113219-dsohsz9y/files/wandb-metadata.json b/wandb/run-20260806_113219-dsohsz9y/files/wandb-metadata.json
new file mode 100644
index 000000000..524b74922
--- /dev/null
+++ b/wandb/run-20260806_113219-dsohsz9y/files/wandb-metadata.json
@@ -0,0 +1,139 @@
+{
+  "os":  "Linux-4.18.0-553.124.1.el8_10.x86_64-x86_64-with-glibc2.28",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-06T16:32:19.879365Z",
+  "args":  [
+    "--ehr-root",
+    "/shared/rsaas/physionet.org/files/mimiciv/2.2",
+    "--note-root",
+    "/shared/rsaas/physionet.org/files/mimic-note",
+    "--cxr-root",
+    "/shared/rsaas/physionet.org/files/MIMIC-CXR",
+    "--cxr-variant",
+    "sunlab",
+    "--cache-dir",
+    "/shared/rsaas/wp14/pyhealth_cache_labs_notes_cxr",
+    "--task",
+    "notes_labs_cxr",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--epochs",
+    "15",
+    "--batch-size",
+    "32",
+    "--lr",
+    "1e-3",
+    "--weight-decay",
+    "1e-5",
+    "--patience",
+    "5",
+    "--num-workers",
+    "4",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/wp14/output",
+    "--pos-weight",
+    "1",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes-cxr",
+    "--wandb-run-name",
+    "labs_notes_cxr_rnn_seed12"
+  ],
+  "program":  "/home/wp14/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "a83f7d6a1dbbc9ff710a88d6027a332e84fde041"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/wp14/PyHealth",
+  "host":  "sunlab-c02.cs.illinois.edu",
+  "executable":  "/home/wp14/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  128,
+  "cpu_count_logical":  255,
+  "gpu":  "NVIDIA RTX 6000 Ada Generation",
+  "gpu_count":  8,
+  "disk":  {
+    "/":  {
+      "total":  "34342961152",
+      "used":  "9715167232"
+    }
+  },
+  "memory":  {
+    "total":  "1081448284160"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-25605dab-2bad-abf9-2f36-f0be6e776096"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fa23eb99-d1b2-4c9c-666c-c70c133af1df"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-39ac7920-1d83-d4cf-c056-920c82599e52"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-3ab5484f-a069-8752-ca10-8ab181aa5b65"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-6a135207-d247-4118-1218-1e2321919d87"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-fdf3835a-4b38-fb01-24cb-f2825806ea1e"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-583da23f-53e9-cf5f-2f63-7115778d21bf"
+    },
+    {
+      "name":  "NVIDIA RTX 6000 Ada Generation",
+      "memoryTotal":  "51527024640",
+      "cudaCores":  18176,
+      "architecture":  "Ada",
+      "uuid":  "GPU-ed1911c1-9405-0826-f8ff-e8ccb5c880c8"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "2e3b1csez9dn1v3fv2wosl25nxbvrw4w"
+}
\ No newline at end of file
diff --git a/wandb/run-20260806_113219-dsohsz9y/run-dsohsz9y.wandb b/wandb/run-20260806_113219-dsohsz9y/run-dsohsz9y.wandb
new file mode 100644
index 000000000..09ce17499
Binary files /dev/null and b/wandb/run-20260806_113219-dsohsz9y/run-dsohsz9y.wandb differ
diff --git a/wandb/run-20260806_231607-4ftnofam/files/config.yaml b/wandb/run-20260806_231607-4ftnofam/files/config.yaml
new file mode 100644
index 000000000..b06c9728e
--- /dev/null
+++ b/wandb/run-20260806_231607-4ftnofam/files/config.yaml
@@ -0,0 +1,216 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            7lrmqxrt80qtcf4rol5v0b483zdsbnyw:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cxr-root
+                    - /home/ubuntu/mimiciv-data/cxr-jpg/2.1.0
+                    - --cxr-variant
+                    - default
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr
+                    - --task
+                    - notes_labs_cxr
+                    - --model
+                    - rnn
+                    - --embedding-dim
+                    - "128"
+                    - --hidden-dim
+                    - "128"
+                    - --rnn-type
+                    - GRU
+                    - --rnn-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes-cxr
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "40225591296"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 9b47d014122d1f6c27ec687a866801c1ea7ea87b
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-2c339b2b-6e2c-c14e-5d5f-e5423a2fffbd
+                host: 192-222-54-167
+                memory:
+                    total: "237490831360"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-06T23:16:07.776027Z"
+                writerId: 7lrmqxrt80qtcf4rol5v0b483zdsbnyw
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr
+cxr_root:
+    value: /home/ubuntu/mimiciv-data/cxr-jpg/2.1.0
+cxr_variant:
+    value: default
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 128
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: rnn
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 2
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs_cxr
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes-cxr
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260806_231607-4ftnofam/files/requirements.txt b/wandb/run-20260806_231607-4ftnofam/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260806_231607-4ftnofam/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260806_231607-4ftnofam/files/wandb-metadata.json b/wandb/run-20260806_231607-4ftnofam/files/wandb-metadata.json
new file mode 100644
index 000000000..22e8c0864
--- /dev/null
+++ b/wandb/run-20260806_231607-4ftnofam/files/wandb-metadata.json
@@ -0,0 +1,91 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-06T23:16:07.776027Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cxr-root",
+    "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0",
+    "--cxr-variant",
+    "default",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr",
+    "--task",
+    "notes_labs_cxr",
+    "--model",
+    "rnn",
+    "--embedding-dim",
+    "128",
+    "--hidden-dim",
+    "128",
+    "--rnn-type",
+    "GRU",
+    "--rnn-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes-cxr",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "9b47d014122d1f6c27ec687a866801c1ea7ea87b"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-54-167",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "40225591296"
+    }
+  },
+  "memory":  {
+    "total":  "237490831360"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-2c339b2b-6e2c-c14e-5d5f-e5423a2fffbd"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "7lrmqxrt80qtcf4rol5v0b483zdsbnyw"
+}
\ No newline at end of file
diff --git a/wandb/run-20260806_231607-4ftnofam/files/wandb-summary.json b/wandb/run-20260806_231607-4ftnofam/files/wandb-summary.json
new file mode 100644
index 000000000..4fdd5f27f
--- /dev/null
+++ b/wandb/run-20260806_231607-4ftnofam/files/wandb-summary.json
@@ -0,0 +1 @@
+{"test_f1":0.06571741511500548,"global_step":90380,"val_pr_auc":0.2304548153067337,"epoch":19,"train_vram_peak_mb":1832.90771484375,"_runtime":17392,"test_roc_auc":0.7534222015833847,"val_roc_auc":0.7501212274868497,"train_vram_allocated_mb":488.4296875,"test_accuracy":0.9528051344472723,"val_accuracy":0.9561223925192276,"_wandb":{"runtime":17392},"pos_weight":1,"val_f1":0.1554845580404686,"_step":20,"test_loss":0.16884233007388832,"_timestamp":1.786075560794198e+09,"skipped_steps":0,"epoch_time_s":738.584,"test_pr_auc":0.23362421417627674,"train_loss":0.1670695467530349,"val_loss":0.16274552233846842}
\ No newline at end of file
diff --git a/wandb/run-20260806_231607-4ftnofam/run-4ftnofam.wandb b/wandb/run-20260806_231607-4ftnofam/run-4ftnofam.wandb
new file mode 100644
index 000000000..1f14728d7
Binary files /dev/null and b/wandb/run-20260806_231607-4ftnofam/run-4ftnofam.wandb differ
diff --git a/wandb/run-20260807_042507-cmlmjn78/files/config.yaml b/wandb/run-20260807_042507-cmlmjn78/files/config.yaml
new file mode 100644
index 000000000..eb48d2041
--- /dev/null
+++ b/wandb/run-20260807_042507-cmlmjn78/files/config.yaml
@@ -0,0 +1,214 @@
+_wandb:
+    value:
+        cli_version: 0.28.0
+        e:
+            p0z3a592esbs9pv0e6ka8u3nvnlo6eo4:
+                args:
+                    - --ehr-root
+                    - /home/ubuntu/mimiciv-data/ehr
+                    - --note-root
+                    - /home/ubuntu/mimiciv-data
+                    - --cxr-root
+                    - /home/ubuntu/mimiciv-data/cxr-jpg/2.1.0
+                    - --cxr-variant
+                    - default
+                    - --cache-dir
+                    - /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr
+                    - --task
+                    - notes_labs_cxr
+                    - --model
+                    - transformer
+                    - --embedding-dim
+                    - "128"
+                    - --heads
+                    - "4"
+                    - --num-layers
+                    - "2"
+                    - --dropout
+                    - "0.1"
+                    - --use-amp
+                    - --amp-dtype
+                    - bf16
+                    - --epochs
+                    - "50"
+                    - --batch-size
+                    - "32"
+                    - --lr
+                    - "0.001"
+                    - --weight-decay
+                    - "1e-05"
+                    - --patience
+                    - "5"
+                    - --pos-weight
+                    - "1"
+                    - --num-workers
+                    - "4"
+                    - --freeze-encoder
+                    - --wandb
+                    - --wandb-project
+                    - pyhealth-multimodal-labs-notes-cxr
+                    - --seed
+                    - "12"
+                    - --output-dir
+                    - /home/ubuntu/output
+                codePath: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                codePathLocal: examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                cpu_count: 13
+                cpu_count_logical: 26
+                cudaVersion: "13.0"
+                disk:
+                    /:
+                        total: "2930253778944"
+                        used: "41540341760"
+                email: williampangbest1@gmail.com
+                executable: /home/ubuntu/miniconda3/envs/pyhealth2/bin/python
+                git:
+                    commit: 9b47d014122d1f6c27ec687a866801c1ea7ea87b
+                    remote: https://github.com/Multimodal-PyHealth/PyHealth
+                gpu: NVIDIA H100 80GB HBM3
+                gpu_count: 1
+                gpu_nvidia:
+                    - architecture: Hopper
+                      cudaCores: 16896
+                      memoryTotal: "85520809984"
+                      name: NVIDIA H100 80GB HBM3
+                      uuid: GPU-2c339b2b-6e2c-c14e-5d5f-e5423a2fffbd
+                host: 192-222-54-167
+                memory:
+                    total: "237490831360"
+                os: Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35
+                program: /home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py
+                python: CPython 3.12.9
+                root: /home/ubuntu/PyHealth
+                startedAt: "2026-08-07T04:25:07.886198Z"
+                writerId: p0z3a592esbs9pv0e6ka8u3nvnlo6eo4
+        m: []
+        python_version: 3.12.9
+        t:
+            "1":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "2":
+                - 1
+                - 5
+                - 11
+                - 41
+                - 49
+                - 53
+                - 54
+                - 71
+                - 98
+            "3":
+                - 2
+                - 13
+                - 15
+                - 16
+                - 61
+            "4": 3.12.9
+            "5": 0.28.0
+            "6": 4.53.3
+            "12": 0.28.0
+            "13": linux-x86_64
+adam_eps:
+    value: null
+amp_dtype:
+    value: bf16
+balanced_ratio:
+    value: 1
+balanced_sampling:
+    value: false
+batch_size:
+    value: 32
+bidirectional:
+    value: false
+bottlenecks_n:
+    value: 4
+cache_dir:
+    value: /home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr
+cxr_root:
+    value: /home/ubuntu/mimiciv-data/cxr-jpg/2.1.0
+cxr_variant:
+    value: default
+dev:
+    value: 0
+device:
+    value: null
+dropout:
+    value: 0.1
+ehr_root:
+    value: /home/ubuntu/mimiciv-data/ehr
+embedding_dim:
+    value: 128
+epochs:
+    value: 50
+freeze_encoder:
+    value: true
+fusion_startidx:
+    value: 1
+heads:
+    value: 4
+hidden_dim:
+    value: 64
+icd_codes:
+    value: false
+include_vitals:
+    value: false
+jamba_mamba_layers:
+    value: 6
+jamba_transformer_layers:
+    value: 2
+lr:
+    value: 0.001
+mamba_conv_kernel:
+    value: 4
+mamba_state_size:
+    value: 16
+max_grad_norm:
+    value: null
+model:
+    value: transformer
+note_root:
+    value: /home/ubuntu/mimiciv-data
+num_layers:
+    value: 2
+num_workers:
+    value: 4
+observation_window_hours:
+    value: 24
+output_dir:
+    value: /home/ubuntu/output
+patience:
+    value: 5
+pos_weight:
+    value: 1
+rnn_layers:
+    value: 1
+rnn_type:
+    value: GRU
+sampling_strategy:
+    value: none
+seed:
+    value: 12
+task:
+    value: notes_labs_cxr
+use_amp:
+    value: true
+wandb:
+    value: true
+wandb_entity:
+    value: null
+wandb_project:
+    value: pyhealth-multimodal-labs-notes-cxr
+wandb_run_name:
+    value: null
+wandb_tags:
+    value: null
+weight_decay:
+    value: 1e-05
diff --git a/wandb/run-20260807_042507-cmlmjn78/files/requirements.txt b/wandb/run-20260807_042507-cmlmjn78/files/requirements.txt
new file mode 100644
index 000000000..6f35326ee
--- /dev/null
+++ b/wandb/run-20260807_042507-cmlmjn78/files/requirements.txt
@@ -0,0 +1,216 @@
+python-dateutil==2.9.0.post0
+triton==3.3.1
+dask==2025.11.0
+typing_utils==0.1.0
+mistune==3.3.2
+pycparser==3.0
+charset-normalizer==3.4.7
+partd==1.4.2
+prompt_toolkit==3.0.52
+packaging==26.0
+nvidia-cuda-runtime-cu12==12.6.77
+webencodings==0.5.1
+s3transfer==0.16.0
+sniffio==1.3.1
+jupyter-lsp==2.3.1
+six==1.17.0
+platformdirs==4.9.6
+decorator==5.2.1
+ipykernel==7.3.0
+brotlicffi==1.2.0.0
+nvidia-curand-cu12==10.3.7.77
+nvidia-cufft-cu12==11.3.0.4
+rpds-py==2026.5.1
+parso==0.8.7
+GitPython==3.1.50
+argon2-cffi-bindings==25.1.0
+nvidia-nvjitlink-cu12==12.6.85
+msgpack==1.1.2
+linformer==0.2.3
+idna==3.11
+psutil==7.2.2
+rfc3339_validator==0.1.4
+pooch==1.9.0
+outdated==0.2.2
+hyper-connections==0.4.9
+comm==0.2.3
+networkx==3.6.1
+nbformat==5.10.4
+anyio==4.14.0
+pexpect==4.9.0
+nbconvert==7.17.1
+soupsieve==2.8.4
+jupyterlab_pygments==0.3.0
+nvidia-cufile-cu12==1.11.1.6
+jsonschema-specifications==2025.9.1
+jupyterlab_widgets==3.0.16
+ptyprocess==0.7.0
+argon2-cffi==25.1.0
+Brotli==1.2.0
+numpy==2.2.6
+notebook_shim==0.2.4
+rdkit==2026.3.1
+pure_eval==0.2.3
+uri-template==1.3.0
+pydantic_core==2.33.2
+pandas==2.3.3
+defusedxml==0.7.1
+scikit-learn==1.7.2
+python-json-logger==4.1.0
+nvidia-cusolver-cu12==11.7.1.2
+pyzmq==27.1.0
+webcolors==25.10.0
+CoLT5-attention==0.11.1
+lark==1.3.1
+ipython==9.14.1
+Jinja2==3.1.6
+fonttools==4.62.1
+ipython_pygments_lexers==1.1.1
+importlib_metadata==9.0.0
+littleutils==0.2.4
+ogb==1.3.6
+regex==2026.4.4
+threadpoolctl==3.6.0
+joblib==1.5.3
+huggingface_hub==0.36.2
+zipp==4.1.0
+polars==1.35.2
+jsonschema==4.26.0
+hpack==4.1.0
+MarkupSafe==3.0.3
+nvidia-cusparse-cu12==12.5.4.2
+tinycss2==1.4.0
+isoduration==20.11.0
+annotated-types==0.7.0
+sympy==1.14.0
+jupyter_builder==1.0.2
+Send2Trash==2.1.0
+jupyter-events==0.12.1
+referencing==0.37.0
+cached-property==1.5.2
+torchvision==0.22.1
+jedi==0.19.2
+jupyter_client==8.9.1
+pillow==12.1.1
+transformers==4.53.3
+babel==2.18.0
+litdata==0.2.61
+tokenizers==0.21.4
+httpx==0.28.1
+linear-attention-transformer==0.19.1
+peft==0.18.1
+jupyter_server==2.20.0
+boto3==1.42.88
+stack_data==0.6.3
+hyperframe==6.1.0
+mne==1.10.2
+beautifulsoup4==4.15.0
+jsonpointer==3.1.1
+nvidia-nvtx-cu12==12.6.77
+wcwidth==0.8.1
+pyarrow==22.0.0
+PyYAML==6.0.3
+wheel==0.46.3
+tblib==3.2.2
+asttokens==3.0.1
+matplotlib-inline==0.2.2
+nvidia-nccl-cu12==2.26.2
+fqdn==1.5.1
+botocore==1.42.88
+rfc3986-validator==0.1.1
+locket==1.0.0
+terminado==0.18.1
+lz4==4.4.5
+cycler==0.12.1
+executing==2.2.1
+traitlets==5.15.1
+ipywidgets==8.1.8
+setuptools==82.0.1
+narwhals==2.13.0
+exceptiongroup==1.3.1
+tzdata==2026.1
+lazy-loader==0.5
+gitdb==4.0.12
+toolz==1.1.0
+fastjsonschema==2.21.2
+nvidia-cuda-cupti-cu12==12.6.80
+arrow==1.4.0
+filelock==3.25.2
+smmap==5.0.3
+notebook==7.6.0
+pip==26.0.1
+nvidia-cusparselt-cu12==0.6.3
+nbclient==0.11.0
+jupyter_core==5.9.1
+typing_extensions==4.15.0
+bokeh==3.9.0
+nest-asyncio2==1.7.2
+rfc3987-syntax==1.1.0
+polars-runtime-32==1.35.2
+product_key_memory==0.3.0
+click==8.3.2
+nvitop==1.7.0
+widgetsnbextension==4.0.15
+safetensors==0.7.0
+websocket-client==1.9.0
+protobuf==6.33.5
+obstore==0.9.2
+PySocks==1.7.1
+sentry-sdk==2.64.0
+scipy==1.17.1
+sortedcontainers==2.4.0
+backports.zstd==1.6.0
+torch-einops-utils==0.0.30
+jmespath==1.1.0
+jupyterlab_server==2.28.0
+fsspec==2026.2.0
+xyzservices==2026.3.0
+distributed==2025.11.0
+certifi==2026.2.25
+pyhealth==2.0.0
+accelerate==1.13.0
+jupyter_server_terminals==0.5.4
+debugpy==1.8.21
+kiwisolver==1.5.0
+tqdm==4.67.3
+einops==0.8.2
+jupyterlab==4.6.0
+Pygments==2.20.0
+nvidia-ml-py==13.590.48
+prometheus_client==0.25.0
+matplotlib==3.10.8
+mpmath==1.3.0
+h11==0.16.0
+overrides==7.7.0
+htcondor==25.11.0
+h2==4.3.0
+tornado==6.5.5
+contourpy==1.3.3
+cffi==1.17.1
+bleach==6.4.0
+urllib3==2.5.0
+attrs==26.1.0
+httpcore==1.0.9
+requests==2.33.1
+async-lru==2.3.0
+axial_positional_embedding==0.3.12
+hf-xet==1.4.3
+tifffile==2026.3.3
+htcondor-cli==25.11.0
+tomli==2.4.1
+nvidia-cublas-cu12==12.6.4.1
+zict==3.0.0
+more-itertools==10.8.0
+typing-inspection==0.4.2
+pandocfilters==1.5.0
+pytz==2026.1.post1
+cloudpickle==3.1.2
+nvidia-cuda-nvrtc-cu12==12.6.77
+wandb==0.28.0
+pyparsing==3.3.2
+torch==2.7.1
+lightning-utilities==0.15.3
+local-attention==1.11.2
+nvidia-cudnn-cu12==9.5.1.17
+pydantic==2.11.10
+json5==0.15.0
diff --git a/wandb/run-20260807_042507-cmlmjn78/files/wandb-metadata.json b/wandb/run-20260807_042507-cmlmjn78/files/wandb-metadata.json
new file mode 100644
index 000000000..13a6e97f8
--- /dev/null
+++ b/wandb/run-20260807_042507-cmlmjn78/files/wandb-metadata.json
@@ -0,0 +1,89 @@
+{
+  "os":  "Linux-6.8.0-1046-nvidia-x86_64-with-glibc2.35",
+  "python":  "CPython 3.12.9",
+  "startedAt":  "2026-08-07T04:25:07.886198Z",
+  "args":  [
+    "--ehr-root",
+    "/home/ubuntu/mimiciv-data/ehr",
+    "--note-root",
+    "/home/ubuntu/mimiciv-data",
+    "--cxr-root",
+    "/home/ubuntu/mimiciv-data/cxr-jpg/2.1.0",
+    "--cxr-variant",
+    "default",
+    "--cache-dir",
+    "/home/ubuntu/mimiciv-data/pyhealth_cache_labs_notes_cxr",
+    "--task",
+    "notes_labs_cxr",
+    "--model",
+    "transformer",
+    "--embedding-dim",
+    "128",
+    "--heads",
+    "4",
+    "--num-layers",
+    "2",
+    "--dropout",
+    "0.1",
+    "--use-amp",
+    "--amp-dtype",
+    "bf16",
+    "--epochs",
+    "50",
+    "--batch-size",
+    "32",
+    "--lr",
+    "0.001",
+    "--weight-decay",
+    "1e-05",
+    "--patience",
+    "5",
+    "--pos-weight",
+    "1",
+    "--num-workers",
+    "4",
+    "--freeze-encoder",
+    "--wandb",
+    "--wandb-project",
+    "pyhealth-multimodal-labs-notes-cxr",
+    "--seed",
+    "12",
+    "--output-dir",
+    "/home/ubuntu/output"
+  ],
+  "program":  "/home/ubuntu/PyHealth/examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePath":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "codePathLocal":  "examples/mortality_prediction/unified_embedding_e2e_mimic4.py",
+  "git":  {
+    "remote":  "https://github.com/Multimodal-PyHealth/PyHealth",
+    "commit":  "9b47d014122d1f6c27ec687a866801c1ea7ea87b"
+  },
+  "email":  "williampangbest1@gmail.com",
+  "root":  "/home/ubuntu/PyHealth",
+  "host":  "192-222-54-167",
+  "executable":  "/home/ubuntu/miniconda3/envs/pyhealth2/bin/python",
+  "cpu_count":  13,
+  "cpu_count_logical":  26,
+  "gpu":  "NVIDIA H100 80GB HBM3",
+  "gpu_count":  1,
+  "disk":  {
+    "/":  {
+      "total":  "2930253778944",
+      "used":  "41540341760"
+    }
+  },
+  "memory":  {
+    "total":  "237490831360"
+  },
+  "gpu_nvidia":  [
+    {
+      "name":  "NVIDIA H100 80GB HBM3",
+      "memoryTotal":  "85520809984",
+      "cudaCores":  16896,
+      "architecture":  "Hopper",
+      "uuid":  "GPU-2c339b2b-6e2c-c14e-5d5f-e5423a2fffbd"
+    }
+  ],
+  "cudaVersion":  "13.0",
+  "writerId":  "p0z3a592esbs9pv0e6ka8u3nvnlo6eo4"
+}
\ No newline at end of file
diff --git a/wandb/run-20260807_042507-cmlmjn78/files/wandb-summary.json b/wandb/run-20260807_042507-cmlmjn78/files/wandb-summary.json
new file mode 100644
index 000000000..d89536f91
--- /dev/null
+++ b/wandb/run-20260807_042507-cmlmjn78/files/wandb-summary.json
@@ -0,0 +1 @@
+{"test_f1":0.2344186046511628,"val_pr_auc":0.27520198449521466,"val_accuracy":0.956841697559896,"train_vram_peak_mb":1828.98388671875,"val_loss":0.15437001430355343,"train_loss":0.15726152462210868,"epoch":13,"test_loss":0.15942153930993735,"epoch_time_s":711.035,"val_roc_auc":0.7900406886843081,"pos_weight":1,"train_vram_allocated_mb":484.5009765625,"test_pr_auc":0.29582848973423276,"test_accuracy":0.9544649773154809,"test_roc_auc":0.80045326232582,"skipped_steps":0,"_wandb":{"runtime":12158},"_step":14,"_timestamp":1.7860888665921793e+09,"val_f1":0.21686746987951808,"_runtime":12158,"global_step":63266}
\ No newline at end of file
diff --git a/wandb/run-20260807_042507-cmlmjn78/run-cmlmjn78.wandb b/wandb/run-20260807_042507-cmlmjn78/run-cmlmjn78.wandb
new file mode 100644
index 000000000..2522c977c
Binary files /dev/null and b/wandb/run-20260807_042507-cmlmjn78/run-cmlmjn78.wandb differ