diff --git a/.gitignore b/.gitignore
index 0d20b64..cddd89c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,45 @@
 *.pyc
+__pycache__/
+*.pyo
+*.pyd
+
+# Virtual environments
+.venv/
+venv/
+env/
+
+# pyenv
+.python-version
+
+# Jupyter
+.ipynb_checkpoints/
+
+# Trained model weights and checkpoints
+*.pt
+*.pth
+*.ckpt
+*.safetensors
+
+# TensorBoard logs
+runs/
+logs/
+
+# Data (potentially large or private datasets)
+data/sudoku/
+data/sudoku-rrn/
+data/clevrtex_full/
+data/clevrtex_outd/
+data/clevrtex_camo/
+data/tetrominoes/
+data/dsprites/
+data/clevr/
+
+# IDE
+.idea/
+.vscode/
+
+# macOS
+.DS_Store
+
+# Environment variables
+.env
diff --git a/FORK.md b/FORK.md
new file mode 100644
index 0000000..fe9c880
--- /dev/null
+++ b/FORK.md
@@ -0,0 +1,170 @@
+# AKOrN – Fork Notes
+
+This is a personal fork of [autonomousvision/akorn](https://github.com/autonomousvision/akorn).
+The upstream repo is the canonical reference for the paper, model architecture, and CLEVRTex experiments.
+This fork adds practical changes needed to run **sudoku evaluation across multiple GPUs** (including Kaggle multi-GPU notebooks) and to develop on macOS.
+
+---
+
+## Motivation
+
+The upstream `eval_sudoku.py` is single-process and single-device.
+Running `K=4096` energy-based voting on the full OOD test set takes several hours on one GPU.
+This fork makes it possible to split that work across N GPUs—each handling a non-overlapping slice of batches—and aggregate the results afterward, without changing the model or the evaluation logic.
+
+Secondary goals: fix a training CSV bug, add `--resume` to `train_sudoku.py`, and make the codebase run on macOS (MPS) or CPU without code changes.
+
+---
+
+## Changes vs upstream
+
+| Area | Change |
+|---|---|
+| `eval_sudoku.py` | Device-agnostic: `torch.load(..., map_location=device)` and `.to(device)` everywhere instead of hardcoded `.cuda()` |
+| `eval_sudoku.py` | `SHARD` / `NSHARD` env vars: each process handles only batch indices where `i % NSHARD == SHARD` |
+| `eval_sudoku.py` | `MAXB` env var: stop after accumulating this many boards (0 = no cap; useful for smoke-tests) |
+| `eval_sudoku.py` | Final print emits `corrects_vote` and `totals` alongside accuracy, so shard outputs can be aggregated |
+| `train_sudoku.py` | `--resume` flag to continue training from a checkpoint |
+| `train_sudoku.py` | CSV logging bug fix (missing flush / truncation issue) |
+| `requirements.txt` | `tensorflow` instead of `tensorflow-cpu` (no macOS ARM wheel for the cpu variant) |
+| `data/download_*.sh` | `curl` instead of `wget` for macOS compatibility |
+
+---
+
+## Multi-GPU evaluation with SHARD / NSHARD / MAXB
+
+### Concept
+
+`NSHARD` is the total number of parallel workers.
+`SHARD` (0-indexed) is the index assigned to this worker.
+Each worker processes only the batches where `batch_index % NSHARD == SHARD`, so there is no overlap and together they cover every batch exactly once.
+
+`MAXB` caps the number of boards a single shard processes (useful for quick sanity checks).
+Set `MAXB=0` (the default) to process the full slice.
+
+### Single-GPU (default, no sharding)
+
+```bash
+python eval_sudoku.py \
+  --data=ood \
+  --model=akorn \
+  --model_path=runs/sudoku_akorn/ema_99.pth \
+  --T=128 --K=4096 --evote_type=sum
+```
+
+Output:
+```
+shard=0/1 corrects_vote=<N> totals=<T> acc=0.XXXX
+```
+
+### Multi-GPU on a single machine (e.g. 4 GPUs)
+
+Run one process per GPU, setting `CUDA_VISIBLE_DEVICES` to pin each process to one device:
+
+```bash
+for SHARD in 0 1 2 3; do
+  CUDA_VISIBLE_DEVICES=$SHARD SHARD=$SHARD NSHARD=4 \
+    python eval_sudoku.py \
+      --data=ood \
+      --model=akorn \
+      --model_path=runs/sudoku_akorn/ema_99.pth \
+      --T=128 --K=4096 --evote_type=sum \
+    > shard_${SHARD}.log 2>&1 &
+done
+wait
+```
+
+Collect results:
+```bash
+grep "shard=" shard_*.log
+# shard=0/4 corrects_vote=1821 totals=2250 acc=0.8093
+# shard=1/4 corrects_vote=1834 totals=2250 acc=0.8151
+# shard=2/4 corrects_vote=1828 totals=2250 acc=0.8124
+# shard=3/4 corrects_vote=1815 totals=2248 acc=0.8072
+```
+
+Aggregate manually (totals may differ by 1 on the last shard due to dataset size):
+```python
+import re, glob
+corrects, totals = 0, 0
+for line in (open(f).read() for f in glob.glob("shard_*.log")):
+    m = re.search(r"corrects_vote=(\d+) totals=(\d+)", line)
+    if m:
+        corrects += int(m.group(1))
+        totals   += int(m.group(2))
+print(f"Overall accuracy: {corrects/totals:.4f}  ({corrects}/{totals})")
+```
+
+### Kaggle notebook (2 GPUs)
+
+Each cell runs on one GPU. Set the env vars at the top of each notebook cell before invoking `eval_sudoku.py`:
+
+**GPU 0 cell:**
+```python
+import os
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["SHARD"]  = "0"
+os.environ["NSHARD"] = "2"
+# os.environ["MAXB"] = "500"  # optional: cap for a quick test
+
+import subprocess
+result = subprocess.run([
+    "python", "eval_sudoku.py",
+    "--data=ood", "--model=akorn",
+    "--model_path=runs/sudoku_akorn/ema_99.pth",
+    "--T=128", "--K=4096", "--evote_type=sum",
+], capture_output=True, text=True)
+print(result.stdout)
+```
+
+**GPU 1 cell** (change `SHARD` to `"1"` and `CUDA_VISIBLE_DEVICES` to `"1"`).
+
+Then aggregate as shown above.
+
+### MAXB – quick sanity check
+
+```bash
+SHARD=0 NSHARD=1 MAXB=200 python eval_sudoku.py \
+  --data=ood --model=akorn \
+  --model_path=runs/sudoku_akorn/ema_99.pth \
+  --T=128 --K=100 --evote_type=sum
+# stops after accumulating 200 boards regardless of dataset size
+```
+
+---
+
+## Training (with resume)
+
+```bash
+# Start
+python train_sudoku.py \
+  --exp_name=sudoku_akorn \
+  --epochs=100 --lr=0.001 --T=16 \
+  --use_omega=True --global_omg=True --init_omg=0.5 --learn_omg=True \
+  --checkpoint_every=10 --eval_freq=10
+
+# Resume from a checkpoint
+python train_sudoku.py \
+  --exp_name=sudoku_akorn \
+  --resume=runs/sudoku_akorn/checkpoint_epoch_50.pth \
+  --epochs=100 --lr=0.001 --T=16 \
+  --use_omega=True --global_omg=True --init_omg=0.5 --learn_omg=True \
+  --checkpoint_every=10 --eval_freq=10
+```
+
+---
+
+## Environment setup
+
+Same as upstream:
+
+```bash
+conda create -n akorn python=3.12 -y
+conda activate akorn
+pip install -r requirements.txt
+```
+
+Data download (macOS-compatible scripts using `curl`):
+```bash
+cd data && bash download_satnet.sh && bash download_rrn.sh && cd ..
+```
\ No newline at end of file
diff --git a/data/download_rrn.sh b/data/download_rrn.sh
index ef2010d..40947d7 100644
--- a/data/download_rrn.sh
+++ b/data/download_rrn.sh
@@ -1,7 +1,6 @@
 # Copied from https://github.com/yilundu/ired_code_release/blob/main/data/download-rrn.sh
 # Original RRN Sodoku data
-wget https://www.dropbox.com/s/rp3hbjs91xiqdgc/sudoku-hard.zip?dl=1
-mv sudoku-hard.zip?dl=1 sudoku-hard.zip
+curl -L "https://www.dropbox.com/s/rp3hbjs91xiqdgc/sudoku-hard.zip?dl=1" -o sudoku-hard.zip
 unzip sudoku-hard.zip
 mv sudoku-hard sudoku-rrn
 rm sudoku-hard.zip
diff --git a/data/download_satnet.sh b/data/download_satnet.sh
index 19e71a2..bd07071 100644
--- a/data/download_satnet.sh
+++ b/data/download_satnet.sh
@@ -1,4 +1,4 @@
 # Copied from https://github.com/yilundu/ired_code_release/blob/main/data/download-satnet.sh
 # Original SAT-Net repo
-wget -cq powei.tw/sudoku.zip && unzip -qq sudoku.zip && rm sudoku.zip
-wget -cq powei.tw/parity.zip && unzip -qq parity.zip && rm parity.zip
\ No newline at end of file
+curl -sL powei.tw/sudoku.zip -o sudoku.zip && unzip -qq sudoku.zip && rm sudoku.zip
+curl -sL powei.tw/parity.zip -o parity.zip && unzip -qq parity.zip && rm parity.zip
\ No newline at end of file
diff --git a/eval_sudoku.py b/eval_sudoku.py
index 10cd1e5..c345416 100644
--- a/eval_sudoku.py
+++ b/eval_sudoku.py
@@ -1,13 +1,9 @@
-import sys, os
+import os
 import torch
-import torch.nn
-import torch.optim
+
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
 import tqdm
-import torchvision
-from torchvision import transforms
 import numpy as np
-from torch.optim.swa_utils import AveragedModel
-import matplotlib.pyplot as plt
 
 from source.data.datasets.sudoku.sudoku import SudokuDataset, HardSudokuDataset
 from source.models.sudoku.knet import SudokuAKOrN
@@ -66,8 +62,9 @@
 
     args = parser.parse_args()
 
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.enable_flash_sdp(enabled=True)
+    if device == "cuda":
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.enable_flash_sdp(enabled=True)
 
     if args.limit_cores_used:
 
@@ -132,7 +129,7 @@ def worker_init_fn(worker_id):
     else:
         raise NotImplementedError
 
-    model = EMA(net).cuda()
+    model = EMA(net).to(device)
     model.load_state_dict(
         torch.load(args.model_path, weights_only=True)["model_state_dict"]
     )
@@ -147,7 +144,15 @@ def worker_init_fn(worker_id):
 
     minimum_chunk = args.minimum_chunk if args.minimum_chunk is not None else K
 
+    SHARD  = int(os.environ.get("SHARD",  0))
+    NSHARD = int(os.environ.get("NSHARD", 1))
+    MAXB   = int(os.environ.get("MAXB",   0))
+
     for i, (X, Y, is_input) in tqdm.tqdm(enumerate(loader)):
+        if NSHARD > 1 and (i % NSHARD) != SHARD:
+            continue
+        if MAXB and totals >= MAXB:
+            break
         B = X.shape[0]
         if args.model == 'akorn' and K > 1:  # Energy-based voting
             for j in range(B):
@@ -159,9 +164,9 @@ def worker_init_fn(worker_id):
                     _Y = Y[j : j + 1].repeat(minimum_chunk, 1, 1, 1)
                     _is_input = is_input[j : j + 1].repeat(minimum_chunk, 1, 1, 1)
                     _X, _Y, _is_input = (
-                        _X.to(torch.int32).cuda(),
-                        _Y.cuda(),
-                        _is_input.cuda(),
+                        _X.to(torch.int32).to(device),
+                        _Y.to(device),
+                        _is_input.to(device),
                     )
 
                     with torch.no_grad():
@@ -189,14 +194,12 @@ def worker_init_fn(worker_id):
                 totals += board_correct_vote.numel()
             
         else:
-            X, Y, is_input = X.to(torch.int32).cuda(), Y.cuda(), is_input.cuda()
+            X, Y, is_input = X.to(torch.int32).to(device), Y.to(device), is_input.to(device)
             with torch.no_grad():
                 pred = model(X, is_input)
             num_blanks, num_corrects, board_correct = compute_board_accuracy(pred, Y, is_input)
             corrects_vote += board_correct.sum().item()
             totals += board_correct.numel()
 
-    # Compute mean and standard deviation across networks
-    accuracy_vote = corrects_vote / totals
-
-    print(f"Vote accuracy: {accuracy_vote:.4f}")
+    accuracy_vote = corrects_vote / totals if totals else 0.0
+    print(f"shard={SHARD}/{NSHARD} corrects_vote={corrects_vote} totals={totals} acc={accuracy_vote:.4f}")
diff --git a/notebooks/kaggle_clevrtex_download.ipynb b/notebooks/kaggle_clevrtex_download.ipynb
new file mode 100644
index 0000000..66b99ab
--- /dev/null
+++ b/notebooks/kaggle_clevrtex_download.ipynb
@@ -0,0 +1,127 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# =============================================================================\n",
+    "# Cell 1 — Config\n",
+    "# CLEVRTex Dataset Download\n",
+    "# Run this notebook ONCE on Kaggle (Internet must be enabled) to download\n",
+    "# CLEVRTex and persist it as a Kaggle Dataset the training notebook can attach.\n",
+    "# Expected runtime: ~30–60 minutes.\n",
+    "# =============================================================================\n",
+    "\n",
+    "import os\n",
+    "import subprocess\n",
+    "import shutil\n",
+    "from pathlib import Path\n",
+    "\n",
+    "OUT_DIR = \"/kaggle/working\"  # clevrtex_full/ subfolder will be created here\n",
+    "\n",
+    "PARTS = [\n",
+    "    \"https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_full_part1.tar.gz\",\n",
+    "    \"https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_full_part2.tar.gz\",\n",
+    "    \"https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_full_part3.tar.gz\",\n",
+    "    \"https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_full_part4.tar.gz\",\n",
+    "    \"https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_full_part5.tar.gz\",\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# =============================================================================\n",
+    "# Cell 2 — Download and extract (one part at a time to stay within 20 GB limit)\n",
+    "# =============================================================================\n",
+    "\n",
+    "for url in PARTS:\n",
+    "    filename = url.split(\"/\")[-1]\n",
+    "    tarball = os.path.join(OUT_DIR, filename)\n",
+    "\n",
+    "    used_gb = shutil.disk_usage(OUT_DIR).used / 1e9\n",
+    "    print(f\"\\n--- Downloading {filename} (disk used before: {used_gb:.2f} GB) ---\")\n",
+    "\n",
+    "    subprocess.run([\"wget\", \"-q\", \"--show-progress\", url, \"-P\", OUT_DIR], check=True)\n",
+    "\n",
+    "    print(f\"Extracting {filename} ...\")\n",
+    "    subprocess.run([\"tar\", \"-xzf\", tarball, \"-C\", OUT_DIR], check=True)\n",
+    "\n",
+    "    os.remove(tarball)\n",
+    "\n",
+    "    used_gb = shutil.disk_usage(OUT_DIR).used / 1e9\n",
+    "    print(f\"Done with {filename}. Disk used after removal: {used_gb:.2f} GB\")\n",
+    "\n",
+    "print(\"\\nAll parts downloaded and extracted.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# =============================================================================\n",
+    "# Cell 3 — Verification\n",
+    "# =============================================================================\n",
+    "\n",
+    "clevrtex_dir = Path(\"/kaggle/working/clevrtex_full\")\n",
+    "\n",
+    "images = list(clevrtex_dir.rglob(\"CLEVRTEX_full_??????.png\"))\n",
+    "print(f\"Image count:  {len(images)}\")\n",
+    "assert len(images) > 0, \"No images found — check that extraction succeeded.\"\n",
+    "\n",
+    "masks = list(clevrtex_dir.rglob(\"*_flat.png\"))\n",
+    "print(f\"Mask count:   {len(masks)}\")\n",
+    "\n",
+    "total_gb = sum(f.stat().st_size for f in clevrtex_dir.rglob(\"*\") if f.is_file()) / 1e9\n",
+    "print(f\"Total size:   {total_gb:.2f} GB\")\n",
+    "\n",
+    "print()\n",
+    "print(\"\\u2713 Dataset ready. Attach this notebook's output to your training notebook as an input dataset.\")\n",
+    "print(\"  Set DATA_ROOT = '/kaggle/input/<this-notebook-slug>' in Cell 1 of the training notebook.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Next Steps\n",
+    "\n",
+    "1. **Save the output.** After this notebook finishes running (\"Save & Run All\"), go to the **Output** tab in the right panel and note the output path.\n",
+    "\n",
+    "2. **Attach to the training notebook.** In `kaggle_clvtex_train.ipynb`, click **Data** (top right) → **Add input** → **Notebook Output Files** → search for this notebook and select its output.\n",
+    "\n",
+    "3. **Set `DATA_ROOT`.** In Cell 1 of `kaggle_clvtex_train.ipynb`, set:\n",
+    "   ```python\n",
+    "   DATA_ROOT = \"/kaggle/input/<this-notebook-slug>\"\n",
+    "   ```\n",
+    "   The CLEVRTex dataset class will automatically look for the `clevrtex_full/` subfolder inside that path.\n",
+    "\n",
+    "4. **Camo / OOD variants.** To evaluate on out-of-distribution variants later, re-run this notebook with additional URLs added to `PARTS`:\n",
+    "   - `https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_outd.tar.gz`\n",
+    "   - `https://thor.robots.ox.ac.uk/datasets/clevrtex/clevrtex_camo.tar.gz`\n",
+    "\n",
+    "   These are **not** needed for the main Table 1 result and can be downloaded separately when needed."
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/notebooks/kaggle_clvtex_train.ipynb b/notebooks/kaggle_clvtex_train.ipynb
new file mode 100644
index 0000000..97b5127
--- /dev/null
+++ b/notebooks/kaggle_clvtex_train.ipynb
@@ -0,0 +1,639 @@
+{
+ "nbformat": 4,
+ "nbformat_minor": 5,
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# AKOrN on CLEVRTex — Multi-Session Training Notebook\n",
+    "\n",
+    "**Target:** Table 1 row \"AKOrN (attn, L=1)\" on CLEVRTex — FG-ARI 75.6, MBO 55.0  \n",
+    "**Platform:** Kaggle T4×2, 12h session cap  \n",
+    "**Repo:** github.com/DimitriBnvl/AKOrN (fork of autonomousvision/akorn)\n",
+    "\n",
+    "---\n",
+    "\n",
+    "## Multi-Session Workflow\n",
+    "\n",
+    "Training 500 epochs (~32–40 hours total) requires multiple Kaggle sessions:\n",
+    "\n",
+    "1. **Session 1 (fresh start):** Leave `PREV_CKPT_INPUT` pointing to a non-existent path (or an empty dataset).  \n",
+    "   Click **Save & Run All** (Commit). Kaggle runs the notebook and saves everything under `/kaggle/working/` as version output.\n",
+    "\n",
+    "2. **Session 2+ (resume):**  \n",
+    "   - Go to your notebook → **Data** → **Add input** → select the *output dataset* of the previous version.  \n",
+    "   - Update `PREV_CKPT_INPUT` to match that dataset's mount path (e.g. `/kaggle/input/akorn-clvtex-ckpts`).  \n",
+    "   - Cell 2 will copy the latest checkpoints into `RUN_DIR` before training starts.  \n",
+    "   - Cell 3 auto-detects the latest checkpoint and sets `FINETUNE_ARG` / `START_EPOCH`.  \n",
+    "   - Cell 5 resumes training from `START_EPOCH` automatically.\n",
+    "\n",
+    "3. **Repeat** until epoch 499 is reached, then run Cell 7 for evaluation.\n",
+    "\n",
+    "### Dataset prerequisites\n",
+    "Upload the 5 CLEVRTex tar.gz parts (`clevrtex_full_part1.tar.gz` … `part5.tar.gz`) to a Kaggle Dataset and set `DATA_ROOT` to its mount path. The `CLEVRTEX` dataset class will look for a `clevrtex_full/` subfolder under that path automatically."
+   ],
+   "id": "markdown-intro",
+   "source": [
+    "# AKOrN on CLEVRTex — Multi-Session Training Notebook\n",
+    "\n",
+    "**Target:** Table 1 row \"AKOrN (attn, L=1)\" on CLEVRTex — FG-ARI 75.6, MBO 55.0  \n",
+    "**Platform:** Kaggle T4×2, 12h session cap  \n",
+    "**Repo:** github.com/DimitriBnvl/AKOrN (fork of autonomousvision/akorn)\n",
+    "\n",
+    "---\n",
+    "\n",
+    "## Multi-Session Workflow\n",
+    "\n",
+    "Training 500 epochs (~32–40 hours total) requires multiple Kaggle sessions:\n",
+    "\n",
+    "1. **Session 1 (fresh start):** Leave `PREV_CKPT_INPUT` pointing to a non-existent path (or an empty dataset).  \n",
+    "   Click **Save & Run All** (Commit). Kaggle runs the notebook and saves everything under `/kaggle/working/` as version output.\n",
+    "\n",
+    "2. **Session 2+ (resume):**  \n",
+    "   - Go to your notebook → **Data** → **Add input** → select the *output dataset* of the previous version.  \n",
+    "   - Update `PREV_CKPT_INPUT` to match that dataset's mount path (e.g. `/kaggle/input/akorn-clvtex-ckpts`).  \n",
+    "   - Cell 2 will copy the latest checkpoints into `RUN_DIR` before training starts.  \n",
+    "   - Cell 3 auto-detects the latest checkpoint and sets `FINETUNE_ARG` / `START_EPOCH`.  \n",
+    "   - Cell 5 resumes training from `START_EPOCH` automatically.\n",
+    "\n",
+    "3. **Repeat** until epoch 499 is reached, then run Cell 7 for evaluation.\n",
+    "\n",
+    "### Dataset prerequisites\n",
+    "Upload the 5 CLEVRTex tar.gz parts (`clevrtex_full_part1.tar.gz` … `part5.tar.gz`) to a Kaggle Dataset and set `DATA_ROOT` to its mount path. The `CLEVRTEX` dataset class will look for a `clevrtex_full/` subfolder under that path automatically."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 1 — Configuration\n",
+    "All paths and hyperparameters in one place. Edit these before committing."
+   ],
+   "id": "markdown-cell1",
+   "source": [
+    "## Cell 1 — Configuration\n",
+    "All paths and hyperparameters in one place. Edit these before committing."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# ── Paths ──────────────────────────────────────────────────────────────────────\n",
+    "REPO_DIR        = \"/kaggle/working/AKOrN\"\n",
+    "DATA_ROOT       = \"/kaggle/input/clevrtex-full\"       # set to your Kaggle Dataset name\n",
+    "PREV_CKPT_INPUT = \"/kaggle/input/akorn-clvtex-ckpts\"  # set to previous version output, or leave absent to start fresh\n",
+    "\n",
+    "# ── Experiment ─────────────────────────────────────────────────────────────────\n",
+    "RUN_NAME = \"clvtex_akorn\"\n",
+    "RUN_DIR  = f\"{REPO_DIR}/runs/{RUN_NAME}\"\n",
+    "\n",
+    "# ── Hyperparameters (match README / Table 1 row) ───────────────────────────────\n",
+    "EPOCHS      = 500\n",
+    "BATCHSIZE   = 256\n",
+    "LR          = 1e-3\n",
+    "CKPT_EVERY  = 50\n",
+    "NUM_GPUS    = 2\n",
+    "\n",
+    "print(\"Configuration loaded.\")\n",
+    "print(f\"  REPO_DIR        = {REPO_DIR}\")\n",
+    "print(f\"  DATA_ROOT       = {DATA_ROOT}\")\n",
+    "print(f\"  PREV_CKPT_INPUT = {PREV_CKPT_INPUT}\")\n",
+    "print(f\"  RUN_DIR         = {RUN_DIR}\")\n",
+    "print(f\"  EPOCHS={EPOCHS}  BATCHSIZE={BATCHSIZE}  LR={LR}  CKPT_EVERY={CKPT_EVERY}  NUM_GPUS={NUM_GPUS}\")"
+   ],
+   "id": "cell-config"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 2 — Setup (idempotent)\n",
+    "Clone repo, install deps, create run dir, and restore checkpoints from a previous session."
+   ],
+   "id": "markdown-cell2",
+   "source": [
+    "## Cell 2 — Setup (idempotent)\n",
+    "Clone repo, install deps, create run dir, and restore checkpoints from a previous session."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import glob\n",
+    "import shutil\n",
+    "import subprocess\n",
+    "\n",
+    "# ── 1. Clone repository (skip if already present) ─────────────────────────────\n",
+    "if not os.path.isdir(REPO_DIR):\n",
+    "    print(\"Cloning AKOrN repository...\")\n",
+    "    subprocess.run(\n",
+    "        [\"git\", \"clone\", \"https://github.com/DimitriBnvl/AKOrN.git\", REPO_DIR],\n",
+    "        check=True,\n",
+    "    )\n",
+    "    print(\"Clone complete.\")\n",
+    "else:\n",
+    "    print(f\"Repository already present at {REPO_DIR}, skipping clone.\")\n",
+    "\n",
+    "# ── 2. Install Python dependencies ────────────────────────────────────────────\n",
+    "print(\"\\nInstalling requirements...\")\n",
+    "subprocess.run(\n",
+    "    [\"pip\", \"install\", \"-q\", \"-r\", os.path.join(REPO_DIR, \"requirements.txt\")],\n",
+    "    check=True,\n",
+    ")\n",
+    "print(\"Requirements installed.\")\n",
+    "\n",
+    "# ── 3. Create run directory ───────────────────────────────────────────────────\n",
+    "os.makedirs(RUN_DIR, exist_ok=True)\n",
+    "print(f\"\\nRun directory ready: {RUN_DIR}\")\n",
+    "\n",
+    "# ── 4. Restore checkpoints from previous session ──────────────────────────────\n",
+    "if os.path.isdir(PREV_CKPT_INPUT):\n",
+    "    prev_ckpts = glob.glob(os.path.join(PREV_CKPT_INPUT, \"**\", \"*.pth\"), recursive=True)\n",
+    "    if prev_ckpts:\n",
+    "        print(f\"\\nFound {len(prev_ckpts)} checkpoint(s) in {PREV_CKPT_INPUT}:\")\n",
+    "        for src in sorted(prev_ckpts):\n",
+    "            dst = os.path.join(RUN_DIR, os.path.basename(src))\n",
+    "            if os.path.exists(dst):\n",
+    "                print(f\"  [skip] {os.path.basename(src)} already in RUN_DIR\")\n",
+    "            else:\n",
+    "                shutil.copy2(src, dst)\n",
+    "                print(f\"  [copy] {os.path.basename(src)}  ({os.path.getsize(src) / 1e6:.1f} MB)\")\n",
+    "    else:\n",
+    "        print(f\"\\nNo .pth files found under {PREV_CKPT_INPUT} — starting fresh.\")\n",
+    "else:\n",
+    "    print(f\"\\nPREV_CKPT_INPUT not found ({PREV_CKPT_INPUT}) — starting fresh.\")\n",
+    "\n",
+    "print(\"\\nSetup complete.\")"
+   ],
+   "id": "cell-setup"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 3 — Verification Gate\n",
+    "Checks GPUs, data, and checkpoint state before committing any compute."
+   ],
+   "id": "markdown-cell3",
+   "source": [
+    "## Cell 3 — Verification Gate\n",
+    "Checks GPUs, data, and checkpoint state before committing any compute."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "import os\n",
+    "import torch\n",
+    "\n",
+    "errors = []\n",
+    "\n",
+    "# ── Check 1: exactly 2 GPUs ───────────────────────────────────────────────────\n",
+    "n_gpu = torch.cuda.device_count()\n",
+    "if n_gpu != 2:\n",
+    "    errors.append(f\"Expected 2 GPUs, found {n_gpu}. Select 'GPU T4 x2' accelerator in notebook settings.\")\n",
+    "\n",
+    "# ── Check 2: both GPUs are T4 (sm_75+, required for Flash SDP) ────────────────\n",
+    "if n_gpu >= 1:\n",
+    "    for i in range(n_gpu):\n",
+    "        name = torch.cuda.get_device_name(i)\n",
+    "        if \"T4\" not in name:\n",
+    "            errors.append(\n",
+    "                f\"GPU {i} is '{name}' — P100/Pascal lacks Flash SDP (sm_75+), use T4 x2.\"\n",
+    "            )\n",
+    "\n",
+    "# ── Check 3: CLEVRTex data is present ────────────────────────────────────────\n",
+    "expected_data_dir = os.path.join(DATA_ROOT, \"clevrtex_full\")\n",
+    "if not os.path.isdir(expected_data_dir):\n",
+    "    errors.append(\n",
+    "        f\"CLEVRTex data not found at '{expected_data_dir}'. \"\n",
+    "        f\"Ensure the Kaggle Dataset is attached and DATA_ROOT='{DATA_ROOT}' is correct.\"\n",
+    "    )\n",
+    "\n",
+    "# ── Check 4 & 5: find latest checkpoint and set resume variables ───────────────\n",
+    "ckpt_pattern = os.path.join(RUN_DIR, \"checkpoint_*.pth\")\n",
+    "existing_ckpts = sorted(glob.glob(ckpt_pattern))\n",
+    "\n",
+    "if existing_ckpts:\n",
+    "    FINETUNE_ARG = existing_ckpts[-1]                            # absolute path\n",
+    "    _meta        = torch.load(FINETUNE_ARG, weights_only=True)\n",
+    "    START_EPOCH  = int(_meta[\"epoch\"]) + 1                      # resume from next epoch\n",
+    "    resume_msg   = f\"Resume from epoch {START_EPOCH} (checkpoint: {os.path.basename(FINETUNE_ARG)})\"\n",
+    "else:\n",
+    "    FINETUNE_ARG = None\n",
+    "    START_EPOCH  = 0\n",
+    "    resume_msg   = \"No checkpoint found — training from scratch (epoch 0)\"\n",
+    "\n",
+    "# ── Print all errors then raise if any ───────────────────────────────────────\n",
+    "if errors:\n",
+    "    print(\"VERIFICATION FAILED:\")\n",
+    "    for e in errors:\n",
+    "        print(f\"  [ERROR] {e}\")\n",
+    "    raise RuntimeError(f\"{len(errors)} verification error(s). See messages above.\")\n",
+    "\n",
+    "# ── Summary ──────────────────────────────────────────────────────────────────\n",
+    "print(\"Verification passed.\")\n",
+    "print()\n",
+    "for i in range(n_gpu):\n",
+    "    props = torch.cuda.get_device_properties(i)\n",
+    "    print(f\"  GPU {i}: {torch.cuda.get_device_name(i)}  \"\n",
+    "          f\"{props.total_memory / 1e9:.1f} GB  sm_{props.major}{props.minor}\")\n",
+    "print(f\"  Data  : {expected_data_dir}\")\n",
+    "print(f\"  {resume_msg}\")\n",
+    "print(f\"  FINETUNE_ARG = {FINETUNE_ARG}\")\n",
+    "print(f\"  START_EPOCH  = {START_EPOCH}\")\n",
+    "print(f\"  Epochs remaining: {EPOCHS - START_EPOCH} / {EPOCHS}\")"
+   ],
+   "id": "cell-verify"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 4 — Timing Probe\n",
+    "Single-GPU forward+backward pass to estimate wall-clock time per epoch and total sessions needed."
+   ],
+   "id": "markdown-cell4",
+   "source": [
+    "## Cell 4 — Timing Probe\n",
+    "Single-GPU forward+backward pass to estimate wall-clock time per epoch and total sessions needed."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "import time\n",
+    "import torch\n",
+    "import torch.nn as nn\n",
+    "\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "os.chdir(REPO_DIR)\n",
+    "\n",
+    "from source.models.objs.knet import AKOrN\n",
+    "\n",
+    "# ── Build model (matches Table 1 / README config) ─────────────────────────────\n",
+    "net = AKOrN(\n",
+    "    4, ch=256, L=1, T=8, gamma=1.0, J=\"attn\",\n",
+    "    use_omega=False, global_omg=False, c_norm=\"gn\",\n",
+    "    psize=8, imsize=128, autorescale=False,\n",
+    "    init_omg=0.01, learn_omg=False,\n",
+    "    maxpool=True, project=True, heads=8,\n",
+    "    use_ro_x=False, no_ro=False, gta=True,\n",
+    ").cuda()\n",
+    "\n",
+    "opt = torch.optim.Adam(net.parameters(), lr=LR)\n",
+    "\n",
+    "# Batch: train_obj.py uses batchsize // num_processes per GPU, then augmentation\n",
+    "# doubles images (2 views per image), so the tensor seen per GPU is:\n",
+    "#   (BATCHSIZE // NUM_GPUS) images × 2 views = (BATCHSIZE // NUM_GPUS * 2) frames\n",
+    "per_gpu_images = BATCHSIZE // NUM_GPUS          # 128\n",
+    "fake = torch.randn(per_gpu_images * 2, 3, 128, 128, device=\"cuda\")  # 2-view batch\n",
+    "\n",
+    "WARMUP_STEPS = 3\n",
+    "TIMED_STEPS  = 20\n",
+    "\n",
+    "torch.backends.cudnn.benchmark = True\n",
+    "torch.backends.cuda.enable_flash_sdp(enabled=True)\n",
+    "\n",
+    "net.train()\n",
+    "\n",
+    "# Warm-up\n",
+    "for _ in range(WARMUP_STEPS):\n",
+    "    out = net(fake)\n",
+    "    # Proxy loss: use out.mean() instead of SimCLR to avoid cross-process gather\n",
+    "    # (SimCLR calls all_gather which requires a real DDP process group).\n",
+    "    # This measures forward+backward compute time only; loss shape doesn't affect timing.\n",
+    "    loss = out.mean()\n",
+    "    opt.zero_grad()\n",
+    "    loss.backward()\n",
+    "    opt.step()\n",
+    "\n",
+    "torch.cuda.synchronize()\n",
+    "t0 = time.perf_counter()\n",
+    "\n",
+    "for _ in range(TIMED_STEPS):\n",
+    "    out = net(fake)\n",
+    "    loss = out.mean()  # proxy loss — see comment above\n",
+    "    opt.zero_grad()\n",
+    "    loss.backward()\n",
+    "    opt.step()\n",
+    "\n",
+    "torch.cuda.synchronize()\n",
+    "t1 = time.perf_counter()\n",
+    "\n",
+    "# ── Compute estimates ─────────────────────────────────────────────────────────\n",
+    "secs_per_step  = (t1 - t0) / TIMED_STEPS\n",
+    "\n",
+    "# Training images: ~40,000 in CLEVRTex train split\n",
+    "TRAIN_IMAGES    = 40_000\n",
+    "STEPS_PER_EPOCH = (TRAIN_IMAGES // NUM_GPUS) // (BATCHSIZE // NUM_GPUS)  # = 156\n",
+    "\n",
+    "# DDP adds communication overhead vs single-GPU timing (~20%)\n",
+    "DDP_OVERHEAD    = 1.2\n",
+    "\n",
+    "secs_per_epoch  = secs_per_step * DDP_OVERHEAD * STEPS_PER_EPOCH\n",
+    "mins_per_epoch  = secs_per_epoch / 60.0\n",
+    "total_hours     = secs_per_epoch * EPOCHS / 3600.0\n",
+    "remaining_hours = secs_per_epoch * (EPOCHS - START_EPOCH) / 3600.0\n",
+    "\n",
+    "SESSION_CAP_H   = 11.5  # Kaggle hard-kills at 12h; leave 30 min margin\n",
+    "sessions_needed = remaining_hours / SESSION_CAP_H\n",
+    "\n",
+    "print(\"Timing probe results\")\n",
+    "print(f\"  Steps per epoch  (formula): {STEPS_PER_EPOCH}\")\n",
+    "print(f\"  s/step  (single GPU, measured): {secs_per_step:.3f}\")\n",
+    "print(f\"  s/step  (2-GPU est. w/ {DDP_OVERHEAD:.1f}x DDP overhead): {secs_per_step * DDP_OVERHEAD:.3f}\")\n",
+    "print(f\"  min/epoch (est.):  {mins_per_epoch:.2f}\")\n",
+    "print(f\"  Total hours for {EPOCHS} ep (est.):     {total_hours:.1f} h\")\n",
+    "print(f\"  Remaining hours from ep {START_EPOCH} (est.): {remaining_hours:.1f} h\")\n",
+    "print(f\"  Sessions needed at {SESSION_CAP_H}h cap: {sessions_needed:.1f}\")\n",
+    "\n",
+    "MAX_SESSIONS = 8\n",
+    "assert sessions_needed < MAX_SESSIONS, (\n",
+    "    f\"Estimated {sessions_needed:.1f} sessions needed exceeds {MAX_SESSIONS}-session sanity limit. \"\n",
+    "    f\"Consider reducing EPOCHS or increasing BATCHSIZE. \"\n",
+    "    f\"Current: EPOCHS={EPOCHS}, START_EPOCH={START_EPOCH}, remaining={remaining_hours:.1f}h.\"\n",
+    ")\n",
+    "\n",
+    "# ── Clean up before training ──────────────────────────────────────────────────\n",
+    "del net, opt, fake\n",
+    "torch.cuda.empty_cache()\n",
+    "print(\"\\nGPU memory released. Ready for training.\")"
+   ],
+   "id": "cell-timing"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 5 — Training\n",
+    "Launch DDP training via `accelerate`. Output is streamed line by line."
+   ],
+   "id": "markdown-cell5",
+   "source": [
+    "## Cell 5 — Training\n",
+    "Launch DDP training via `accelerate`. Output is streamed line by line."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "import subprocess\n",
+    "\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "os.chdir(REPO_DIR)  # train_obj.py writes 'runs/<exp_name>/' relative to cwd\n",
+    "\n",
+    "# ── Build accelerate launch command ───────────────────────────────────────────\n",
+    "cmd = [\n",
+    "    \"accelerate\", \"launch\",\n",
+    "    \"--multi_gpu\",\n",
+    "    f\"--num_processes={NUM_GPUS}\",\n",
+    "    \"train_obj.py\",\n",
+    "    f\"--exp_name={RUN_NAME}\",\n",
+    "    f\"--data_root={DATA_ROOT}\",   # CLEVRTEX class appends clevrtex_full/ subfolder automatically\n",
+    "    \"--model=akorn\",\n",
+    "    \"--data=clevrtex_full\",\n",
+    "    \"--J=attn\",\n",
+    "    \"--L=1\",\n",
+    "    f\"--epochs={EPOCHS}\",\n",
+    "    f\"--batchsize={BATCHSIZE}\",\n",
+    "    f\"--lr={LR}\",\n",
+    "    f\"--checkpoint_every={CKPT_EVERY}\",\n",
+    "]\n",
+    "\n",
+    "# Append resume checkpoint if one was detected in Cell 3\n",
+    "if FINETUNE_ARG is not None:\n",
+    "    cmd.append(f\"--finetune={FINETUNE_ARG}\")\n",
+    "\n",
+    "print(\"Training command:\")\n",
+    "print(\" \".join(cmd))\n",
+    "print(f\"\\nResuming from epoch: {START_EPOCH} / {EPOCHS}\")\n",
+    "print(\"-\" * 72)\n",
+    "\n",
+    "# ── Stream subprocess output line by line ─────────────────────────────────────\n",
+    "proc = subprocess.Popen(\n",
+    "    cmd,\n",
+    "    stdout=subprocess.PIPE,\n",
+    "    stderr=subprocess.STDOUT,\n",
+    "    text=True,\n",
+    "    bufsize=1,\n",
+    "    cwd=REPO_DIR,\n",
+    ")\n",
+    "\n",
+    "for line in proc.stdout:\n",
+    "    print(line, end=\"\", flush=True)\n",
+    "\n",
+    "proc.wait()\n",
+    "\n",
+    "if proc.returncode != 0:\n",
+    "    raise subprocess.CalledProcessError(proc.returncode, cmd)\n",
+    "\n",
+    "print(\"-\" * 72)\n",
+    "print(\"Training completed successfully.\")"
+   ],
+   "id": "cell-train"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 6 — Checkpoint Status\n",
+    "Inventory of all saved checkpoints and EMA files. Informational only — never raises."
+   ],
+   "id": "markdown-cell6",
+   "source": [
+    "## Cell 6 — Checkpoint Status\n",
+    "Inventory of all saved checkpoints and EMA files. Informational only — never raises."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import glob\n",
+    "import os\n",
+    "import torch\n",
+    "\n",
+    "print(f\"Checkpoint inventory: {RUN_DIR}\")\n",
+    "print(\"-\" * 72)\n",
+    "\n",
+    "ckpt_files = sorted(glob.glob(os.path.join(RUN_DIR, \"checkpoint_*.pth\")))\n",
+    "ema_files  = sorted(glob.glob(os.path.join(RUN_DIR, \"ema_*.pth\")))\n",
+    "all_pth    = sorted(glob.glob(os.path.join(RUN_DIR, \"*.pth\")))\n",
+    "\n",
+    "total_bytes = 0\n",
+    "\n",
+    "if not all_pth:\n",
+    "    print(\"  (no .pth files found)\")\n",
+    "else:\n",
+    "    for fpath in all_pth:\n",
+    "        size_mb = os.path.getsize(fpath) / 1e6\n",
+    "        total_bytes += os.path.getsize(fpath)\n",
+    "        try:\n",
+    "            meta  = torch.load(fpath, weights_only=True)\n",
+    "            epoch = meta.get(\"epoch\", \"n/a\")\n",
+    "        except Exception as exc:\n",
+    "            epoch = f\"(load error: {exc})\"\n",
+    "        fname = os.path.basename(fpath)\n",
+    "        print(f\"  {fname:<30s}  epoch={epoch}  {size_mb:.1f} MB\")\n",
+    "\n",
+    "print(\"-\" * 72)\n",
+    "print(f\"  Total: {len(all_pth)} files  |  {total_bytes / 1e9:.2f} GB\")\n",
+    "print()\n",
+    "\n",
+    "# Quick summary\n",
+    "if ckpt_files:\n",
+    "    latest_meta = torch.load(ckpt_files[-1], weights_only=True)\n",
+    "    latest_ep   = latest_meta.get(\"epoch\", \"?\")\n",
+    "    print(f\"Latest checkpoint: epoch {latest_ep} ({os.path.basename(ckpt_files[-1])})\")\n",
+    "    print(f\"Epochs complete: {int(latest_ep)+1} / {EPOCHS}\")\n",
+    "    remaining = EPOCHS - int(latest_ep) - 1\n",
+    "    print(f\"Epochs remaining: {remaining}\")\n",
+    "else:\n",
+    "    print(\"No checkpoint_*.pth files saved yet.\")"
+   ],
+   "id": "cell-ckpt-status"
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Cell 7 — Evaluation\n",
+    "Run `eval_obj.py` on the final EMA model to obtain FG-ARI and MBO on the CLEVRTex test set.\n",
+    "\n",
+    "**Target (Table 1, AKOrN attn L=1):** FG-ARI = 75.6, MBO = 55.0"
+   ],
+   "id": "markdown-cell7",
+   "source": [
+    "## Cell 7 — Evaluation\n",
+    "Run `eval_obj.py` on the final EMA model to obtain FG-ARI and MBO on the CLEVRTex test set.\n",
+    "\n",
+    "**Target (Table 1, AKOrN attn L=1):** FG-ARI = 75.6, MBO = 55.0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import os\n",
+    "import subprocess\n",
+    "\n",
+    "sys.path.insert(0, REPO_DIR)\n",
+    "os.chdir(REPO_DIR)\n",
+    "\n",
+    "# ── Target checkpoint (EMA at final epoch 499) ────────────────────────────────\n",
+    "EVAL_MODEL_PATH = os.path.join(RUN_DIR, \"ema_499.pth\")\n",
+    "\n",
+    "assert os.path.isfile(EVAL_MODEL_PATH), (\n",
+    "    f\"Evaluation checkpoint not found: {EVAL_MODEL_PATH}\\n\"\n",
+    "    \"Training may not be complete, or previous sessions' EMA files need to be restored. \"\n",
+    "    \"Attach the final session's output dataset and re-run Cell 2 to copy checkpoints.\"\n",
+    ")\n",
+    "\n",
+    "print(f\"Evaluating: {EVAL_MODEL_PATH}\")\n",
+    "print(\"Target: FG-ARI = 75.6, MBO = 55.0  (Table 1, AKOrN attn L=1 on CLEVRTex)\")\n",
+    "print(\"-\" * 72)\n",
+    "\n",
+    "# ── Main eval: CLEVRTex full test set (FG-ARI + MBO) ─────────────────────────\n",
+    "eval_cmd = [\n",
+    "    \"python\", os.path.join(REPO_DIR, \"eval_obj.py\"),\n",
+    "    f\"--data_root={DATA_ROOT}\",   # CLEVRTEX class appends clevrtex_full/ automatically\n",
+    "    \"--model=akorn\",\n",
+    "    \"--data=clevrtex_full\",\n",
+    "    \"--J=attn\",\n",
+    "    \"--L=1\",\n",
+    "    f\"--model_path={EVAL_MODEL_PATH}\",\n",
+    "    \"--model_imsize=128\",\n",
+    "]\n",
+    "\n",
+    "print(\"Eval command:\")\n",
+    "print(\" \".join(eval_cmd))\n",
+    "print()\n",
+    "\n",
+    "proc = subprocess.Popen(\n",
+    "    eval_cmd,\n",
+    "    stdout=subprocess.PIPE,\n",
+    "    stderr=subprocess.STDOUT,\n",
+    "    text=True,\n",
+    "    bufsize=1,\n",
+    "    cwd=REPO_DIR,\n",
+    ")\n",
+    "\n",
+    "for line in proc.stdout:\n",
+    "    print(line, end=\"\", flush=True)\n",
+    "\n",
+    "proc.wait()\n",
+    "\n",
+    "if proc.returncode != 0:\n",
+    "    raise subprocess.CalledProcessError(proc.returncode, eval_cmd)\n",
+    "\n",
+    "print(\"-\" * 72)\n",
+    "print(\"Evaluation complete.\")\n",
+    "\n",
+    "# ── Optional OOD and CAMO evaluations (commented out) ─────────────────────────\n",
+    "# Uncomment to evaluate on out-of-distribution splits:\n",
+    "#\n",
+    "# EVAL_OOD_CMD = [\n",
+    "#     \"python\", os.path.join(REPO_DIR, \"eval_obj.py\"),\n",
+    "#     f\"--data_root={DATA_ROOT}\",\n",
+    "#     \"--model=akorn\", \"--data=clevrtex_outd\",\n",
+    "#     \"--J=attn\", \"--L=1\",\n",
+    "#     f\"--model_path={EVAL_MODEL_PATH}\",\n",
+    "#     \"--model_imsize=128\",\n",
+    "# ]\n",
+    "#\n",
+    "# EVAL_CAMO_CMD = [\n",
+    "#     \"python\", os.path.join(REPO_DIR, \"eval_obj.py\"),\n",
+    "#     f\"--data_root={DATA_ROOT}\",\n",
+    "#     \"--model=akorn\", \"--data=clevrtex_camo\",\n",
+    "#     \"--J=attn\", \"--L=1\",\n",
+    "#     f\"--model_path={EVAL_MODEL_PATH}\",\n",
+    "#     \"--model_imsize=128\",\n",
+    "# ]\n",
+    "#\n",
+    "# for name, cmd in [(\"OOD\", EVAL_OOD_CMD), (\"CAMO\", EVAL_CAMO_CMD)]:\n",
+    "#     print(f\"\\n--- {name} eval ---\")\n",
+    "#     subprocess.run(cmd, check=True, cwd=REPO_DIR)"
+   ],
+   "id": "cell-eval"
+  }
+ ]
+}
diff --git a/requirements.txt b/requirements.txt
index 35477f7..fbdd36d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -8,7 +8,7 @@ tensorboard
 tqdm
 argparse
 git+https://github.com/fra31/auto-attack
-tensorflow-cpu
+tensorflow
 ema_pytorch
 accelerate
 scipy
diff --git a/run_sudoku.sh b/run_sudoku.sh
new file mode 100644
index 0000000..2b6752e
--- /dev/null
+++ b/run_sudoku.sh
@@ -0,0 +1,15 @@
+# Training Command for Mac
+python train_sudoku.py --exp_name=sudoku_akorn --eval_freq=10 --epochs=100 --model=akorn --lr=0.001 --T=16 --use_omega=True --global_omg=True --init_omg=0.5 --learn_omg=True --num_workers=0 --checkpoint_every=10
+
+# Training Command for Linux / CUDA (with checkpoints)
+python train_sudoku.py --exp_name=sudoku_akorn  --eval_freq=10 --epochs=100 --model=akorn --lr=0.001 --T=16 --use_omega=True --global_omg=True --init_omg=0.5 --learn_omg=True --checkpoint_every=10
+
+# Evaluation
+export data=ood # id or ood
+
+# Inference with test-time extension of the Kuramoto updates. (Accuracy: 51.7%)
+python eval_sudoku.py --data=${data} --model=akorn --model_path=runs/sudoku_akorn/ema_99.pth --T=128
+# Test-time extension and energy-based voting (Accuracy: 81.6%)
+python eval_sudoku.py --data=${data} --model=akorn --model_path=runs/sudoku_akorn/ema_99.pth --T=128 --K=100 --evote_type=sum
+# Number of random samples increased from 100 to 4096 (best results) (Accuracy: 89.5%)
+python eval_sudoku.py --data=${data} --model=akorn --model_path=runs/sudoku_akorn/ema_99.pth --T=128 --K=4096 --evote_type=sum
diff --git a/source/data/datasets/sudoku/sudoku.py b/source/data/datasets/sudoku/sudoku.py
index bcc2b0a..5c26a67 100644
--- a/source/data/datasets/sudoku/sudoku.py
+++ b/source/data/datasets/sudoku/sudoku.py
@@ -21,7 +21,7 @@ def load_rrn_dataset(data_dir, split):
     split_to_filename = {"train": "train.csv", "val": "valid.csv", "test": "test.csv"}
 
     filename = osp.join(data_dir, split_to_filename[split])
-    df = pd.read_csv(filename, header=None)
+    df = pd.read_csv(filename, header=None, dtype=str)
 
     def str2onehot(x):
         x = np.array(list(map(int, x)), dtype="int64")
diff --git a/source/layers/klayer.py b/source/layers/klayer.py
index 1bd9a4b..ef0056e 100644
--- a/source/layers/klayer.py
+++ b/source/layers/klayer.py
@@ -26,9 +26,6 @@ def __init__(self, n, ch, init_omg=0.1, global_omg=False, learn_omg=True):
         self.ch = ch
         self.global_omg = global_omg
 
-        if not learn_omg:
-            print("Not learning omega")
-
         if n % 2 != 0:
             # n is odd
             raise NotImplementedError
diff --git a/train_obj.py b/train_obj.py
index e399f41..99e2b17 100644
--- a/train_obj.py
+++ b/train_obj.py
@@ -362,7 +362,9 @@ def train(net, ema, opt, scheduler, loader, epoch):
 
     scheduler = LinearWarmupScheduler(optimizer, warmup_iters=args.warmup_iters)
 
-    for epoch in range(0, args.epochs):
+    start_epoch = torch.load(args.finetune)["epoch"] + 1 if args.finetune else 0
+
+    for epoch in range(start_epoch, args.epochs):
         total_loss = train(net, ema, optimizer, scheduler, ssloader, epoch)
         if (epoch + 1) % args.checkpoint_every == 0:
             if accelerator.is_main_process:
@@ -372,6 +374,7 @@ def train(net, ema, opt, scheduler, loader, epoch):
                     epoch,
                     total_loss,
                     checkpoint_dir=jobdir,
+                    max_checkpoints=2,
                 )
                 save_model(ema, epoch, checkpoint_dir=jobdir, prefix="ema")
     if accelerator.is_main_process:
diff --git a/train_sudoku.py b/train_sudoku.py
index 740fd32..25a9d71 100644
--- a/train_sudoku.py
+++ b/train_sudoku.py
@@ -3,6 +3,8 @@
 import tqdm
 import argparse
 
+device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+
 from source.models.sudoku.transformer import SudokuTransformer
 
 from source.training_utils import save_checkpoint, save_model
@@ -76,14 +78,16 @@ def apply_threshold(model, threshold):
     parser.add_argument("--init_omg", type=float, default=0.1)
     parser.add_argument("--nl", type=str2bool, default=True)
 
+    parser.add_argument("--resume", action="store_true", help="resume from latest checkpoint")
     parser.add_argument("--speed_test", action="store_true")
 
     args = parser.parse_args()
 
     print("Exp name: ", args.exp_name)
 
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.enable_flash_sdp(enabled=True)
+    if device == "cuda":
+        torch.backends.cudnn.benchmark = True
+        torch.backends.cuda.enable_flash_sdp(enabled=True)
     
     if args.seed is not None:
         import random
@@ -128,7 +132,7 @@ def compute_acc(net, loader):
         correct_input = 0
         total_input = 0
         for X, Y, is_input in loader:
-            X, Y, is_input = X.to(torch.int32).cuda(), Y.cuda(), is_input.cuda()
+            X, Y, is_input = X.to(torch.int32).to(device), Y.to(device), is_input.to(device)
 
             with torch.no_grad():
                 out = net(X, is_input)
@@ -178,7 +182,7 @@ def compute_acc(net, loader):
     else:
         raise NotImplementedError
 
-    net.cuda()
+    net.to(device)
 
     total_params = sum(p.numel() for p in net.parameters() if p.requires_grad)
     print(f"Total number of parameters: {total_params}")
@@ -187,6 +191,28 @@ def compute_acc(net, loader):
 
     ema = EMA(net, beta=args.beta, update_every=10, update_after_step=100)
 
+    start_epoch = 0
+    if args.resume:
+        checkpoints = [
+            f for f in os.listdir(jobdir)
+            if f.startswith("checkpoint_") and f.endswith(".pth")
+        ]
+        if checkpoints:
+            checkpoints.sort(key=lambda f: int(f.split("_")[1].split(".")[0]))
+            latest = checkpoints[-1]
+            latest_epoch = int(latest.split("_")[1].split(".")[0])
+            ckpt = torch.load(os.path.join(jobdir, latest), map_location=device, weights_only=True)
+            net.load_state_dict(ckpt["model_state_dict"])
+            optimizer.load_state_dict(ckpt["optimizer_state_dict"])
+            ema_path = os.path.join(jobdir, f"ema_{latest_epoch}.pth")
+            if os.path.exists(ema_path):
+                ema_ckpt = torch.load(ema_path, map_location=device, weights_only=True)
+                ema.load_state_dict(ema_ckpt["model_state_dict"])
+            start_epoch = latest_epoch + 1
+            print(f"Resumed from epoch {latest_epoch}, continuing from epoch {start_epoch}")
+        else:
+            print("No checkpoints found, starting from scratch")
+
     criterion = torch.nn.CrossEntropyLoss(reduction="none")
 
     # Measure speed
@@ -194,14 +220,14 @@ def compute_acc(net, loader):
         it_sp = 0
         time_per_iter = []
         import numpy as np
-        
-    for epoch in range(args.epochs):
+
+    for epoch in range(start_epoch, args.epochs):
         total_loss = 0
 
         for X, Y, is_input in tqdm.tqdm(trainloader):
             net.train()
             ema.train()
-            X, Y, is_input = X.to(torch.int32).cuda(), Y.cuda(), is_input.cuda()
+            X, Y, is_input = X.to(torch.int32).to(device), Y.to(device), is_input.to(device)
 
             if args.speed_test:
                 start = torch.cuda.Event(enable_timing=True)