From 47486cf079a6d07571f07183f3ae7178fec3c1c3 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Tue, 18 Feb 2025 15:19:35 +0000
Subject: [PATCH 1/7] MI325 - rocm 6.2

---
 config/base.yaml            | 58 ++++++++++++++++++-------------------
 milabench/_version.py       |  6 ++--
 scripts/article/run_rocm.sh | 42 +++++++++++++++++----------
 3 files changed, 58 insertions(+), 48 deletions(-)

diff --git a/config/base.yaml b/config/base.yaml
index 82ff3a9a..da6b8ef6 100644
--- a/config/base.yaml
+++ b/config/base.yaml
@@ -148,7 +148,7 @@ _accelerate_opt:
     --dataset_config_name: "wikitext-103-v1"
     --dataset_rev: "b08601e"
     --validation_split_percentage: 5
-    --per_gpu_batch_size: 1
+    --per_gpu_batch_size: "auto_batch(1)"
     --cpus_per_gpu: "auto({n_worker}, 8)"
     --cache: "{milabench_cache}"
 
@@ -229,7 +229,7 @@ resnet50-noio:
   
   argv:
     --model: resnet50
-    --batch-size: batch_resize(256)
+    --batch-size: "auto_batch(256)"
     --loader: synthetic_fixed
     --optim: channel_last
 
@@ -244,7 +244,7 @@ resnet152-ddp-gpus:
   
   argv:
     --model: resnet152
-    --batch-size: 256
+    --batch-size: "auto_batch(256)"
     --num-workers: "auto({n_worker}, 8)"
     --loader: torch
 
@@ -294,7 +294,7 @@ regnet_y_128gf:
     - monogpu
   argv:
     --model: regnet_y_128gf
-    --batch-size: 64
+    --batch-size: "auto_batch(64)"
 
 _bert-base:
   inherits: _hf
@@ -308,7 +308,7 @@ _bert-base:
     - monogpu
   argv:
     --model: "Bert"
-    --batch-size: 32
+    --batch-size: "auto_batch(32)"
   voir:
     options:
       stop: 30
@@ -344,7 +344,7 @@ t5:
     - monogpu
   argv:
     --model: "T5"
-    --batch-size: 16
+    --batch-size: "auto_batch(16)"
 
 reformer:
   inherits: _hf
@@ -357,7 +357,7 @@ reformer:
     - monogpu
   argv:
     --model: "Reformer"
-    --batch-size: 32
+    --batch-size: "auto_batch(32)"
 
 whisper:
   inherits: _hf
@@ -368,7 +368,7 @@ whisper:
     - monogpu
   argv:
     --model: "Whisper"
-    --batch-size: 64
+    --batch-size: "auto_batch(64)"
 
 focalnet:
   inherits: _timm
@@ -400,7 +400,7 @@ brax:
     n: 1
   argv:
     --episode-length: 20
-    --batch-size: 1024
+    --batch-size: "auto_batch(1024)"
     --num-minibatches: 32
     --num-envs: 8192
 
@@ -415,7 +415,7 @@ _diffusion:
 
   argv:
     --num_epochs: 5
-    --batch_size: 32
+    --batch_size: "auto_batch(32)"
     --num_workers: "auto({n_worker}, 8)"
     --cache: "{milabench_cache}"
 
@@ -459,7 +459,7 @@ _lightning:
     --loader: pytorch
     --data: "{milabench_data}/FakeImageNet"
     --model: resnet152
-    --batch-size: 256
+    --batch-size: "auto_batch(256)"
 
 lightning:
   inherits: _lightning
@@ -504,7 +504,7 @@ dinov2-giant-single:
     --config-file: "{benchmark_folder}/src/dinov2/configs/train/vitg14.yaml"
     # THOSE NEED TO BE LAST
     train.dataset_path=ImageNet:split=TRAIN:root={milabench_data}/FakeImageNet:extra={milabench_data}/FakeImageNet: true
-    train.batch_size_per_gpu=32: true
+    train.batch_size_per_gpu=auto_batch(32): true
     train.saveckp_freq=100: true
     train.num_workers=10: true
 
@@ -519,7 +519,7 @@ dinov2-giant-gpus:
     --config-file: "{benchmark_folder}/src/dinov2/configs/train/vitg14.yaml"
     # THOSE NEED TO BE LAST
     train.dataset_path=ImageNet:split=TRAIN:root={milabench_data}/FakeImageNet:extra={milabench_data}/FakeImageNet: true
-    train.batch_size_per_gpu=32: true
+    train.batch_size_per_gpu=auto_batch(32): true
     train.saveckp_freq=100: true
     train.num_workers=10: true
 
@@ -535,7 +535,7 @@ dinov2-giant-nodes:
     --config-file: "{benchmark_folder}/src/dinov2/configs/train/vitg14.yaml"
     # THOSE NEED TO BE LAST
     train.dataset_path=ImageNet:split=TRAIN:root={milabench_data}/FakeImageNet:extra={milabench_data}/FakeImageNet: true
-    train.batch_size_per_gpu=12: true
+    train.batch_size_per_gpu=auto_batch(12): true
     train.saveckp_freq=100: true
     train.num_workers=10: true
 
@@ -576,7 +576,7 @@ llm-lora-single:
     checkpointer.checkpoint_dir={milabench_data}/llama3_8B: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
-    batch_size=8: true
+    batch_size=auto_batch(8): true
     gradient_accumulation_steps=8: true
     device={device_name}: true
 
@@ -602,7 +602,7 @@ llm-lora-ddp-gpus:
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
-    batch_size=8: true
+    batch_size=auto_batch(8): true
     gradient_accumulation_steps=8: true
     device={device_name}: true
 
@@ -629,7 +629,7 @@ llm-lora-ddp-nodes:
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
-    batch_size=8: true
+    batch_size=auto_batch(8): true
     gradient_accumulation_steps=8: true
     device={device_name}: true
   num_machines: 2
@@ -656,7 +656,7 @@ llm-lora-mp-gpus:
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
-    batch_size=8: true
+    batch_size=auto_batch(8): true
     gradient_accumulation_steps=1: true
     device={device_name}: true
   
@@ -685,7 +685,7 @@ llm-full-mp-gpus:
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
     safetensors=true: true
-    batch_size=2: true
+    batch_size=auto_batch(2): true
     gradient_accumulation_steps=1: true
     device={device_name}: true
   
@@ -710,7 +710,7 @@ llm-full-mp-nodes:
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
     safetensors=true: true
-    batch_size=2: true
+    batch_size=auto_batch(2): true
     gradient_accumulation_steps=1: true
     device={device_name}: true
   
@@ -766,7 +766,7 @@ pna:
   argv:
     --model: 'PNA'
     --num-samples: 100000
-    --batch-size: 4096
+    --batch-size: "auto_batch(4096)" 
     --num-workers: "auto({n_worker}, 0)"
 
 dimenet:
@@ -775,7 +775,7 @@ dimenet:
     --model: 'DimeNet'
     --num-samples: 100000
     --use3d: True
-    --batch-size: 16
+    --batch-size: "auto_batch(16)"
     --num-workers: "auto({n_worker}, 0)"
 
 recursiongfn:
@@ -789,7 +789,7 @@ recursiongfn:
     method: per_gpu
 
   argv:
-    --batch_size: 128
+    --batch_size: "auto_batch(128)"
     --num_workers: "auto({n_worker}, 8)"
     --num_steps: 100
     --layer_width: 128
@@ -806,7 +806,7 @@ torchatari:
     - monogpu
     - gym
   argv:
-    --num-minibatches: 16
+    --num-minibatches: "auto_batch(16)"
     --update-epochs: 4
     --num-steps: 128
     --num-envs: auto({cpu_per_gpu}, 128)
@@ -823,7 +823,7 @@ _llava:
   tags:
     - llm
   argv:
-    --batch_size: 1
+    --batch_size: "auto_batch(1)"
     --num_workers: "auto({n_worker}, 4)"
     --gradient_accumulation_steps: 1
 
@@ -834,7 +834,7 @@ llava-single:
   plan:
     method: per_gpu
   argv:
-    --batch_size: 1
+    --batch_size: "auto_batch(1)"
     --num_workers: "auto({n_worker}, 4)"
     --gradient_accumulation_steps: 1
 
@@ -846,7 +846,7 @@ llava-gpus:
     method: njobs
     n: 1
   argv:
-    --batch_size: 1
+    --batch_size: "auto_batch(1)"
     --num_workers: "auto({n_worker}, 4)"
     --gradient_accumulation_steps: 1
 
@@ -865,7 +865,7 @@ _rlhf:
     --output_dir: "{milabench_extra}/output"
     --model_name_or_path: EleutherAI/pythia-1b-deduped
 
-    --per_device_train_batch_size: 64
+    --per_device_train_batch_size: "auto_batch(64)"
     --logging_strategy: "no"
     --log_level: "critical"
     --bf16: true
@@ -894,7 +894,7 @@ _vjepa:
   tags:
     - video
   argv:
-    --batch_size: 24
+    --batch_size: "auto_batch(24)"
     --num_workers: "auto({n_worker}, 12)"
     --dataset: "{milabench_data}/FakeVideo/video_metainfo.csv"
     --output: "{milabench_extra}"
diff --git a/milabench/_version.py b/milabench/_version.py
index 71d24764..d109cfd9 100644
--- a/milabench/_version.py
+++ b/milabench/_version.py
@@ -1,5 +1,5 @@
 """This file is generated, do not modify"""
 
-__tag__ = "v1.0.0-24-g776e3e1"
-__commit__ = "776e3e1bba8be548c22fd4a811092e7e00128689"
-__date__ = "2025-01-16 15:14:33 -0500"
+__tag__ = "v1.0.0-50-g3249edf"
+__commit__ = "3249edfd466030ba2d7a612ac8b8102f6253a5a7"
+__date__ = "2025-02-06 10:51:24 -0500"
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index 0fc2bf16..4b4221d8 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -5,6 +5,10 @@ set -ex
 # sudo usermod -a -G render,video $LOGNAME
 # sudo chmod u+s /opt/rocm-6.2.2/lib/llvm/bin/amdgpu-arch
 
+
+# sudo apt-get install g++-12 libstdc++-12-dev
+
+
 export MILABENCH_GPU_ARCH=rocm
 export MILABENCH_WORDIR="$(pwd)/$MILABENCH_GPU_ARCH"
 export ROCM_PATH="/opt/rocm"
@@ -84,25 +88,28 @@ install_prepare() {
         # Override/add package to the benchmark venv here
         #
         which pip
+        export MAX_JOBS=24
+        # https://github.com/ROCm/jax/releases/download/rocm-jax-v0.5.0/jaxlib-0.5.0-cp310-cp310-manylinux_2_28_x86_64.whl
 
         # https://github.com/ROCm/jax/releases/tag/rocm-jaxlib-v0.4.30
-        pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp310-cp310-manylinux2014_x86_64.whl
-        pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.g
+        # pip install https://github.com/ROCm/jax/releases/download/rocm-jaxlib-v0.4.30/jaxlib-0.4.30+rocm611-cp310-cp310-manylinux2014_x86_64.whl
+        # pip install https://github.com/ROCm/jax/archive/refs/tags/rocm-jaxlib-v0.4.30.tar.gz
 
-        pip uninstall torch_cluster torch_scatter torch_sparse -y
-        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_cluster.git
-        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_scatter.git
-        FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
+        # pip uninstall torch_cluster torch_scatter torch_sparse -y
+        # FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_cluster.git
+        # FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_scatter.git
+        # FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
 
         # takes forever to compile
         # https://github.com/ROCm/xformers
-        pip uninstall xformers
-        pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
-        # pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
+        pip uninstall xformers -y
+        # pip install "torch<2.6" --index-url https://download.pytorch.org/whl/rocm6.2
+        # pip install xformers==0.0.29 --index-url https://download.pytorch.org/whl/rocm6.2
+        pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
         # pip install -v -U --no-build-isolation --no-deps git+https://github.com/facebookresearch/xformers.git
         # pip install xformers -U --index-url https://download.pytorch.org/whl/rocm6.1
 
-        pip uninstall flash-attention
+        pip uninstall -y flash-attention
         pip install -v -U --no-build-isolation --use-pep517 --no-deps git+https://github.com/ROCm/flash-attention.git 
         pip uninstall pynvml nvidia-ml-py -y
 
@@ -123,20 +130,23 @@ else
     . $MILABENCH_WORDIR/env/bin/activate
 fi
 
-(
-    . $BENCHMARK_VENV/bin/activate
-    pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
-)
+# (
+#     # . $BENCHMARK_VENV/bin/activate
+#     # pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
+# )
 
 # milabench install $ARGS --system $MILABENCH_WORDIR/system.yaml
 
-# milabench prepare $ARGS --system $MILABENCH_WORDIR/system.yaml
+# milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
+# install_prepare
 
 #
 #   Run the benchmakrs
-milabench run $ARGS --system $MILABENCH_WORDIR/system.yaml
+milabench run $ARGS # --system $MILABENCH_WORDIR/system.yaml
 
 
 #
 #   Display report
 milabench report --runs $MILABENCH_WORDIR/results/runs
+
+# rocm
\ No newline at end of file

From 5440826e4cdd300ce0a5c051d03bb6233342d954 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Wed, 26 Feb 2025 19:26:11 +0000
Subject: [PATCH 2/7] Update ROCm to run on 2.6

---
 constraints/extra/torch.rocm.txt |  6 ---
 constraints/rocm.txt             |  2 +-
 scripts/article/run_rocm.sh      | 87 ++++++++++++++------------------
 3 files changed, 39 insertions(+), 56 deletions(-)

diff --git a/constraints/extra/torch.rocm.txt b/constraints/extra/torch.rocm.txt
index 81763330..e69de29b 100644
--- a/constraints/extra/torch.rocm.txt
+++ b/constraints/extra/torch.rocm.txt
@@ -1,6 +0,0 @@
-# No jax only a container for it
-
-
-
-# https://github.com/ROCm/xformers
-# pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
\ No newline at end of file
diff --git a/constraints/rocm.txt b/constraints/rocm.txt
index aeaf8c6f..5c7283c4 100644
--- a/constraints/rocm.txt
+++ b/constraints/rocm.txt
@@ -1,4 +1,4 @@
---extra-index-url https://download.pytorch.org/whl/rocm6.2
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 #
 #
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index 79b6b308..7551b97c 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -4,8 +4,6 @@ set -ex
 
 # sudo usermod -a -G render,video $LOGNAME
 # sudo chmod u+s /opt/rocm-6.2.2/lib/llvm/bin/amdgpu-arch
-
-
 # sudo apt-get install g++-12 libstdc++-12-dev
 
 
@@ -23,23 +21,12 @@ else
     export MILABENCH_CONFIG="$MILABENCH_SOURCE/config/standard.yaml"
 fi
 
-
 export GPU="$(/opt/rocm/lib/llvm/bin/amdgpu-arch | head -n 1)"
 export TORCH_ROCM_ARCH_LIST="$GPU"
 export ROCM_TARGETS="$GPU"
 export PYTORCH_ROCM_ARCH="$GPU"
+export MAX_JOBS=24
 
-if [ -z "${MILABENCH_SOURCE}" ]; then
-    export MILABENCH_CONFIG="$MILABENCH_WORDIR/milabench/config/standard.yaml"
-else
-    export MILABENCH_CONFIG="$MILABENCH_SOURCE/config/standard.yaml"
-fi
-
-
-export GPU="$(/opt/rocm/lib/llvm/bin/amdgpu-arch | head -n 1)"
-export TORCH_ROCM_ARCH_LIST="$GPU"
-export ROCM_TARGETS="$GPU"
-export PYTORCH_ROCM_ARCH="$GPU"
 
 ARGS="$@"
 
@@ -51,8 +38,6 @@ export XLA_FLAGS=--xla_gpu_autotune_level=0
 install_jax() {
     (
         . $BENCHMARK_VENV/bin/activate
-    
-        export MAX_JOBS=24
 
         # Jax 0.4.30
         # https://github.com/ROCm/jax/releases/tag/rocm-jaxlib-v0.4.30
@@ -80,37 +65,9 @@ install_graph() {
     FORCE_ONLY_CUDA=1 pip install -U -v --use-pep517 --no-build-isolation git+https://github.com/rusty1s/pytorch_sparse.git
 }
 
-install_prepare() {
-    mkdir -p $MILABENCH_WORDIR
-    cd $MILABENCH_WORDIR
 
-    virtualenv $MILABENCH_WORDIR/env
-
-    if [ -z "${MILABENCH_SOURCE}" ]; then
-        if [ ! -d "$MILABENCH_WORDIR/milabench" ]; then
-            git clone https://github.com/mila-iqia/milabench.git -b rocm
-        fi
-        export MILABENCH_SOURCE="$MILABENCH_WORDIR/milabench"
-    fi
-
-    . $MILABENCH_WORDIR/env/bin/activate
-    pip install -e $MILABENCH_SOURCE
-
-    
-    #
-    # Install milabench's benchmarks in their venv
-    #
-    # pip install torch --index-url https://download.pytorch.org/whl/rocm6.1
-    # milabench pin --variant rocm --from-scratch $ARGS 
-    milabench install $ARGS 
-
-    #
-    # Override/add package to milabench venv here
-    #
-    which pip
-    pip uninstall pynvml
-
-    (
+install_override() {
+  (
         . $BENCHMARK_VENV/bin/activate
 
         pip install ninja
@@ -133,19 +90,51 @@ install_prepare() {
 
         # takes forever to compile
         # https://github.com/ROCm/xformers
-        pip uninstall xformers -y
+        # pip uninstall xformers -y
         # pip install "torch<2.6" --index-url https://download.pytorch.org/whl/rocm6.2
         # pip install xformers==0.0.29 --index-url https://download.pytorch.org/whl/rocm6.2
-        pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
+        # pip install -v -U --no-build-isolation --no-deps git+https://github.com/ROCm/xformers.git@develop#egg=xformers
         # pip install -v -U --no-build-isolation --no-deps git+https://github.com/facebookresearch/xformers.git
         # pip install xformers -U --index-url https://download.pytorch.org/whl/rocm6.1
 
         pip uninstall -y flash-attention
         pip install -v -U --no-build-isolation --use-pep517 --no-deps git+https://github.com/ROCm/flash-attention.git 
         pip uninstall pynvml nvidia-ml-py -y
-
         pip install einops
     )
+}
+
+install_prepare() {
+    mkdir -p $MILABENCH_WORDIR
+    cd $MILABENCH_WORDIR
+
+    virtualenv $MILABENCH_WORDIR/env
+
+    if [ -z "${MILABENCH_SOURCE}" ]; then
+        if [ ! -d "$MILABENCH_WORDIR/milabench" ]; then
+            git clone https://github.com/mila-iqia/milabench.git -b rocm
+        fi
+        export MILABENCH_SOURCE="$MILABENCH_WORDIR/milabench"
+    fi
+
+    . $MILABENCH_WORDIR/env/bin/activate
+    pip install -e $MILABENCH_SOURCE
+
+    
+    #
+    # Install milabench's benchmarks in their venv
+    #
+    # pip install torch --index-url https://download.pytorch.org/whl/rocm6.1
+    milabench pin --variant rocm --from-scratch $ARGS 
+    milabench install $ARGS 
+
+    #
+    # Override/add package to milabench venv here
+    #
+    which pip
+    pip uninstall pynvml
+
+    install_override
 
     pip uninstall pynvml nvidia-ml-py -y
     #

From cd0eab6110d5b09d53d71e47bfe2dcf3a3e5bb3f Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Wed, 26 Feb 2025 19:59:23 +0000
Subject: [PATCH 3/7] update to pytorch 2.6 + rocm 6.2.4

---
 .pin/constraints-rocm-torch.txt               | 999 ++++++++++++++++++
 benchmarks/brax/requirements.rocm.txt         | 213 ++--
 benchmarks/diffusion/requirements.rocm.txt    | 187 ++--
 benchmarks/dinov2/requirements.rocm.txt       | 117 +-
 benchmarks/flops/requirements.rocm.txt        |  91 +-
 benchmarks/geo_gnn/requirements-pre.rocm.txt  |  33 +-
 benchmarks/geo_gnn/requirements.rocm.txt      | 151 ++-
 benchmarks/huggingface/requirements.rocm.txt  | 111 +-
 benchmarks/lightning/requirements.rocm.txt    | 131 ++-
 benchmarks/llama/requirements.rocm.txt        | 171 ++-
 benchmarks/llava/requirements.rocm.txt        | 169 ++-
 benchmarks/llm/requirements.rocm.txt          | 218 ++--
 benchmarks/purejaxrl/requirements.rocm.txt    | 375 +++----
 benchmarks/recursiongfn/requirements.rocm.txt | 278 ++---
 benchmarks/rlhf/requirements.rocm.txt         | 198 ++--
 benchmarks/timm/requirements.rocm.txt         | 107 +-
 benchmarks/torchatari/requirements.rocm.txt   | 157 +--
 benchmarks/torchvision/requirements.rocm.txt  |  91 +-
 .../torchvision_ddp/requirements.rocm.txt     |  91 +-
 benchmarks/vjepa/requirements.rocm.txt        | 147 ++-
 benchmate/benchmate/warden.py                 |   8 +-
 milabench/_version.py                         |   6 +-
 scripts/article/run_rocm.sh                   |   7 +-
 23 files changed, 2521 insertions(+), 1535 deletions(-)
 create mode 100644 .pin/constraints-rocm-torch.txt

diff --git a/.pin/constraints-rocm-torch.txt b/.pin/constraints-rocm-torch.txt
new file mode 100644
index 00000000..cab2f08a
--- /dev/null
+++ b/.pin/constraints-rocm-torch.txt
@@ -0,0 +1,999 @@
+#
+# This file is autogenerated by pip-compile with Python 3.10
+# by the following command:
+#
+#    pip-compile --output-file=.pin/constraints-rocm-torch.txt .pin/tmp-constraints.txt benchmarks/brax/requirements.in benchmarks/diffusion/requirements.in benchmarks/dinov2/requirements.in benchmarks/flops/requirements.in benchmarks/geo_gnn/requirements-pre.in benchmarks/geo_gnn/requirements.in benchmarks/huggingface/requirements.in benchmarks/lightning/requirements.in benchmarks/llama/requirements.in benchmarks/llava/requirements.in benchmarks/llm/requirements.in benchmarks/purejaxrl/requirements.in benchmarks/recursiongfn/requirements.in benchmarks/rlhf/requirements.in benchmarks/timm/requirements.in benchmarks/torchatari/requirements.in benchmarks/torchvision/requirements.in benchmarks/torchvision_ddp/requirements.in benchmarks/vjepa/requirements.in constraints/extra/torch.rocm.txt
+#
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
+
+absl-py==2.1.0
+    # via
+    #   brax
+    #   chex
+    #   distrax
+    #   dm-env
+    #   dm-tree
+    #   ml-collections
+    #   mujoco
+    #   mujoco-mjx
+    #   optax
+    #   orbax-checkpoint
+    #   rlax
+    #   tensorboard
+    #   tensorflow-probability
+accelerate==1.4.0
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/rlhf/requirements.in
+    #   diffusers
+    #   trl
+aiohappyeyeballs==2.4.6
+    # via aiohttp
+aiohttp==3.11.13
+    # via
+    #   datasets
+    #   fsspec
+    #   torch-geometric
+aiosignal==1.3.2
+    # via aiohttp
+annotated-types==0.7.0
+    # via pydantic
+antlr4-python3-runtime==4.9.3
+    # via omegaconf
+appdirs==1.4.4
+    # via cantilever
+argklass==1.4.4
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+astroid==3.3.8
+    # via pylint
+asttokens==3.0.0
+    # via giving
+async-timeout==5.0.1
+    # via aiohttp
+attrs==25.1.0
+    # via
+    #   aiohttp
+    #   dm-tree
+beartype==0.20.0
+    # via -r benchmarks/vjepa/requirements.in
+black==25.1.0
+    # via navix
+blinker==1.9.0
+    # via flask
+blobfile==3.0.0
+    # via
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
+    #   torchtune
+blosc2==2.7.1
+    # via tables
+botorch==0.13.0
+    # via -r benchmarks/recursiongfn/requirements.in
+braceexpand==0.1.7
+    # via
+    #   -r benchmarks/vjepa/requirements.in
+    #   webdataset
+brax==0.12.1
+    # via
+    #   -r benchmarks/brax/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+cantilever==0.1.0
+    # via -r benchmarks/torchatari/requirements.in
+certifi==2025.1.31
+    # via
+    #   requests
+    #   sentry-sdk
+charset-normalizer==3.4.1
+    # via requests
+chex==0.1.89
+    # via
+    #   distrax
+    #   evosax
+    #   flashbax
+    #   gymnax
+    #   optax
+    #   rlax
+click==8.1.8
+    # via
+    #   black
+    #   flask
+    #   wandb
+cloudpickle==3.1.1
+    # via
+    #   gym
+    #   gymnasium
+    #   submitit
+    #   tensorflow-probability
+codefind==0.1.7
+    # via ptera
+contourpy==1.3.1
+    # via matplotlib
+cvxopt==1.3.2
+    # via -r benchmarks/recursiongfn/requirements.in
+cycler==0.12.1
+    # via matplotlib
+datasets==3.3.2
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/llama/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/rlhf/requirements.in
+    #   torchtune
+    #   trl
+decorator==5.2.1
+    # via tensorflow-probability
+decord==0.6.0
+    # via -r benchmarks/vjepa/requirements.in
+diffusers[torch]==0.32.2
+    # via -r benchmarks/diffusion/requirements.in
+dill==0.3.8
+    # via
+    #   datasets
+    #   multiprocess
+    #   pylint
+distrax==0.1.5
+    # via
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   rlax
+dm-env==1.6
+    # via
+    #   brax
+    #   envpool
+    #   rlax
+dm-tree==0.1.9
+    # via
+    #   dm-env
+    #   tensorflow-probability
+docker-pycreds==0.4.0
+    # via wandb
+docstring-parser==0.16
+    # via tyro
+dotmap==1.3.30
+    # via evosax
+einops==0.8.1
+    # via
+    #   -r benchmarks/rlhf/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+envpool==0.8.4
+    # via -r benchmarks/torchatari/requirements.in
+etils[epath,epy]==1.12.0
+    # via
+    #   brax
+    #   mujoco
+    #   mujoco-mjx
+    #   optax
+    #   orbax-checkpoint
+evosax==0.1.6
+    # via -r benchmarks/purejaxrl/requirements.in
+exceptiongroup==1.2.2
+    # via pytest
+executing==2.2.0
+    # via varname
+fairscale==0.4.13
+    # via
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
+    #   -r benchmarks/llama/requirements.in
+    #   -r benchmarks/llm/requirements.in
+farama-notifications==0.0.4
+    # via gymnasium
+filelock==3.17.0
+    # via
+    #   blobfile
+    #   datasets
+    #   diffusers
+    #   huggingface-hub
+    #   torch
+    #   transformers
+fire==0.7.0
+    # via
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
+    #   -r benchmarks/llama/requirements.in
+flake8==7.1.2
+    # via navix
+flashbax==0.1.2
+    # via -r benchmarks/purejaxrl/requirements.in
+flask==3.1.0
+    # via
+    #   brax
+    #   flask-cors
+flask-cors==5.0.1
+    # via brax
+flax==0.10.3
+    # via
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   brax
+    #   evosax
+    #   flashbax
+    #   gymnax
+    #   navix
+fonttools==4.56.0
+    # via matplotlib
+frozenlist==1.5.0
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec[http]==2024.12.0
+    # via
+    #   datasets
+    #   etils
+    #   huggingface-hub
+    #   lightning
+    #   pytorch-lightning
+    #   torch
+    #   torch-geometric
+fvcore==0.1.5.post20221221
+    # via -r benchmarks/dinov2/requirements.in
+gast==0.6.0
+    # via tensorflow-probability
+gitdb==4.0.12
+    # via gitpython
+gitpython==3.1.44
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   wandb
+giving==0.4.3
+    # via
+    #   ptera
+    #   voir
+glfw==2.8.0
+    # via mujoco
+gpytorch==1.14
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   botorch
+grpcio==1.70.0
+    # via
+    #   brax
+    #   tensorboard
+gym==0.26.2
+    # via
+    #   -r benchmarks/torchatari/requirements.in
+    #   brax
+    #   envpool
+    #   gymnax
+gym-notices==0.0.8
+    # via gym
+gymnasium==1.1.0
+    # via
+    #   envpool
+    #   gymnax
+gymnax==0.0.8
+    # via
+    #   -c /home/testroot/milabench/constraints/rocm.txt
+    #   -r benchmarks/purejaxrl/requirements.in
+hf-transfer==0.1.9
+    # via huggingface-hub
+hjson==3.1.0
+    # via argklass
+huggingface-hub[hf-transfer]==0.29.1
+    # via
+    #   -r benchmarks/timm/requirements.in
+    #   accelerate
+    #   datasets
+    #   diffusers
+    #   timm
+    #   tokenizers
+    #   torchtune
+    #   transformers
+humanize==4.12.1
+    # via orbax-checkpoint
+idna==3.10
+    # via
+    #   requests
+    #   yarl
+importlib-metadata==8.6.1
+    # via diffusers
+importlib-resources==6.5.2
+    # via
+    #   argklass
+    #   cantilever
+    #   etils
+    #   torchcompat
+iniconfig==2.0.0
+    # via pytest
+iopath==0.1.10
+    # via
+    #   -r benchmarks/dinov2/requirements.in
+    #   fvcore
+isort==6.0.0
+    # via pylint
+itsdangerous==2.2.0
+    # via flask
+jax==0.5.1
+    # via
+    #   -r benchmarks/brax/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   brax
+    #   chex
+    #   distrax
+    #   evosax
+    #   flashbax
+    #   flax
+    #   gymnax
+    #   jaxopt
+    #   mujoco-mjx
+    #   optax
+    #   orbax-checkpoint
+    #   rlax
+jaxlib==0.5.1
+    # via
+    #   brax
+    #   chex
+    #   distrax
+    #   evosax
+    #   flashbax
+    #   gymnax
+    #   jax
+    #   jaxopt
+    #   mujoco-mjx
+    #   optax
+    #   rlax
+jaxopt==0.8.3
+    # via brax
+jaxtyping==0.2.38
+    # via
+    #   gpytorch
+    #   linear-operator
+jinja2==3.1.5
+    # via
+    #   brax
+    #   flask
+    #   torch
+    #   torch-geometric
+joblib==1.4.2
+    # via scikit-learn
+kagglehub==0.3.10
+    # via torchtune
+kiwisolver==1.4.8
+    # via matplotlib
+lightning==2.5.0.post0
+    # via -r benchmarks/lightning/requirements.in
+lightning-utilities==0.12.0
+    # via
+    #   lightning
+    #   pytorch-lightning
+    #   torchmetrics
+linear-operator==0.6
+    # via
+    #   botorch
+    #   gpytorch
+lxml==5.3.1
+    # via blobfile
+markdown==3.7
+    # via tensorboard
+markdown-it-py==3.0.0
+    # via rich
+markupsafe==3.0.2
+    # via
+    #   jinja2
+    #   werkzeug
+matplotlib==3.10.0
+    # via
+    #   evosax
+    #   gymnax
+    #   seaborn
+mccabe==0.7.0
+    # via
+    #   flake8
+    #   pylint
+mdurl==0.1.2
+    # via markdown-it-py
+ml-collections==1.0.0
+    # via brax
+ml-dtypes==0.5.1
+    # via
+    #   jax
+    #   jaxlib
+    #   tensorstore
+mpmath==1.3.0
+    # via
+    #   gpytorch
+    #   linear-operator
+    #   sympy
+msgpack==1.1.0
+    # via
+    #   blosc2
+    #   flax
+    #   orbax-checkpoint
+mujoco==3.2.7
+    # via
+    #   brax
+    #   mujoco-mjx
+mujoco-mjx==3.2.7
+    # via brax
+multidict==6.1.0
+    # via
+    #   aiohttp
+    #   yarl
+multipledispatch==1.0.0
+    # via botorch
+multiprocess==0.70.16
+    # via datasets
+mypy-extensions==1.0.0
+    # via
+    #   black
+    #   typing-inspect
+navix==0.7.0
+    # via -r benchmarks/purejaxrl/requirements.in
+ndindex==1.9.2
+    # via blosc2
+nest-asyncio==1.6.0
+    # via orbax-checkpoint
+networkx==3.4.2
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   torch
+numexpr==2.10.2
+    # via
+    #   blosc2
+    #   tables
+numpy==1.26.4
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   -r benchmarks/torchatari/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+    #   accelerate
+    #   blosc2
+    #   brax
+    #   chex
+    #   contourpy
+    #   datasets
+    #   decord
+    #   diffusers
+    #   distrax
+    #   dm-env
+    #   dm-tree
+    #   envpool
+    #   evosax
+    #   fairscale
+    #   flashbax
+    #   fvcore
+    #   gym
+    #   gymnasium
+    #   jax
+    #   jaxlib
+    #   jaxopt
+    #   matplotlib
+    #   ml-dtypes
+    #   mujoco
+    #   navix
+    #   numexpr
+    #   opencv-python
+    #   optax
+    #   orbax-checkpoint
+    #   pandas
+    #   pyro-ppl
+    #   rdkit
+    #   rlax
+    #   scikit-learn
+    #   scipy
+    #   seaborn
+    #   tables
+    #   tensorboard
+    #   tensorboardx
+    #   tensorflow-probability
+    #   tensorstore
+    #   torch-geometric
+    #   torchmetrics
+    #   torchtune
+    #   torchvision
+    #   transformers
+    #   treescope
+    #   trimesh
+    #   webdataset
+    #   xformers
+nvidia-ml-py==12.570.86
+    # via voir
+omegaconf==2.3.0
+    # via
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   torchtune
+    #   voir
+opencv-python==4.11.0.86
+    # via -r benchmarks/vjepa/requirements.in
+opt-einsum==3.4.0
+    # via
+    #   jax
+    #   pyro-ppl
+optax==0.2.4
+    # via
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   brax
+    #   flax
+optree==0.14.0
+    # via envpool
+orbax-checkpoint==0.11.6
+    # via
+    #   brax
+    #   flax
+ovld==0.3.9
+    # via voir
+packaging==24.2
+    # via
+    #   accelerate
+    #   black
+    #   datasets
+    #   envpool
+    #   huggingface-hub
+    #   kagglehub
+    #   lightning
+    #   lightning-utilities
+    #   matplotlib
+    #   pytest
+    #   pytorch-lightning
+    #   setuptools-scm
+    #   tables
+    #   tensorboard
+    #   tensorboardx
+    #   torchmetrics
+    #   transformers
+pandas==2.2.3
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+    #   datasets
+    #   seaborn
+pathspec==0.12.1
+    # via black
+pillow==11.1.0
+    # via
+    #   -r benchmarks/huggingface/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   brax
+    #   diffusers
+    #   fvcore
+    #   matplotlib
+    #   navix
+    #   rdkit
+    #   torchtune
+    #   torchvision
+platformdirs==4.3.6
+    # via
+    #   black
+    #   pylint
+    #   wandb
+pluggy==1.5.0
+    # via pytest
+portalocker==3.1.1
+    # via iopath
+propcache==0.3.0
+    # via
+    #   aiohttp
+    #   yarl
+protobuf==5.29.3
+    # via
+    #   orbax-checkpoint
+    #   tensorboard
+    #   tensorboardx
+    #   wandb
+psutil==5.9.8
+    # via
+    #   accelerate
+    #   torch-geometric
+    #   torchtune
+    #   voir
+    #   wandb
+ptera==1.4.1
+    # via voir
+py-cpuinfo==9.0.0
+    # via
+    #   blosc2
+    #   tables
+pyarrow==19.0.1
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   datasets
+pycodestyle==2.12.1
+    # via flake8
+pycryptodomex==3.21.0
+    # via blobfile
+pydantic==2.10.6
+    # via wandb
+pydantic-core==2.27.2
+    # via pydantic
+pyflakes==3.2.0
+    # via flake8
+pygments==2.19.1
+    # via rich
+pylint==3.3.4
+    # via navix
+pyopengl==3.1.9
+    # via mujoco
+pyparsing==3.2.1
+    # via
+    #   matplotlib
+    #   torch-geometric
+pyre-extensions==0.0.32
+    # via botorch
+pyro-api==0.1.2
+    # via pyro-ppl
+pyro-ppl==1.9.1
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   botorch
+pytest==8.3.4
+    # via navix
+python-dateutil==2.9.0.post0
+    # via
+    #   matplotlib
+    #   pandas
+pytinyrenderer==0.0.14
+    # via brax
+pytorch-lightning==2.5.0.post0
+    # via lightning
+pytorch-triton-rocm==3.2.0
+    # via torch
+pytz==2025.1
+    # via pandas
+pyyaml==6.0.2
+    # via
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/timm/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+    #   accelerate
+    #   datasets
+    #   evosax
+    #   flax
+    #   fvcore
+    #   gymnax
+    #   huggingface-hub
+    #   kagglehub
+    #   lightning
+    #   ml-collections
+    #   omegaconf
+    #   orbax-checkpoint
+    #   pytorch-lightning
+    #   timm
+    #   transformers
+    #   wandb
+    #   webdataset
+    #   yacs
+rdkit==2024.9.5
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+reactivex==4.0.4
+    # via giving
+regex==2024.11.6
+    # via
+    #   diffusers
+    #   tiktoken
+    #   transformers
+requests==2.32.3
+    # via
+    #   datasets
+    #   diffusers
+    #   huggingface-hub
+    #   kagglehub
+    #   tiktoken
+    #   torch-geometric
+    #   transformers
+    #   wandb
+rich==13.9.4
+    # via
+    #   flax
+    #   trl
+    #   tyro
+    #   voir
+rlax==0.1.6
+    # via navix
+safetensors==0.5.3
+    # via
+    #   -r benchmarks/timm/requirements.in
+    #   accelerate
+    #   diffusers
+    #   timm
+    #   torchtune
+    #   transformers
+scikit-learn==1.6.1
+    # via gpytorch
+scipy==1.15.2
+    # via
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   botorch
+    #   brax
+    #   gpytorch
+    #   jax
+    #   jaxlib
+    #   jaxopt
+    #   linear-operator
+    #   mujoco-mjx
+    #   scikit-learn
+    #   torch-cluster
+    #   torch-sparse
+seaborn==0.13.2
+    # via gymnax
+sentencepiece==0.2.0
+    # via
+    #   -r benchmarks/llama/requirements.in
+    #   torchtune
+sentry-sdk==2.22.0
+    # via wandb
+setproctitle==1.3.5
+    # via wandb
+setuptools-scm==8.2.0
+    # via navix
+shtab==1.7.1
+    # via tyro
+simplejson==3.20.1
+    # via orbax-checkpoint
+six==1.17.0
+    # via
+    #   docker-pycreds
+    #   ml-collections
+    #   python-dateutil
+    #   tensorboard
+    #   tensorflow-probability
+smmap==5.0.2
+    # via gitdb
+submitit==1.5.2
+    # via
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+sympy==1.13.1
+    # via torch
+tables==3.10.1
+    # via -r benchmarks/recursiongfn/requirements.in
+tabulate==0.9.0
+    # via fvcore
+tensorboard==2.19.0
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   -r benchmarks/torchatari/requirements.in
+tensorboard-data-server==0.7.2
+    # via tensorboard
+tensorboardx==2.6.2.2
+    # via brax
+tensorflow-probability==0.25.0
+    # via distrax
+tensorstore==0.1.72
+    # via
+    #   flashbax
+    #   flax
+    #   orbax-checkpoint
+termcolor==2.5.0
+    # via
+    #   fire
+    #   fvcore
+threadpoolctl==3.5.0
+    # via
+    #   botorch
+    #   scikit-learn
+tiktoken==0.9.0
+    # via torchtune
+timm==1.0.15
+    # via -r benchmarks/vjepa/requirements.in
+tokenizers==0.21.0
+    # via transformers
+tomli==2.2.1
+    # via
+    #   black
+    #   pylint
+    #   pytest
+    #   setuptools-scm
+tomlkit==0.13.2
+    # via pylint
+toolz==1.0.0
+    # via chex
+torch==2.6.0+rocm6.2.4
+    # via
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
+    #   -r benchmarks/brax/requirements.in
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/flops/requirements.in
+    #   -r benchmarks/geo_gnn/requirements-pre.in
+    #   -r benchmarks/huggingface/requirements.in
+    #   -r benchmarks/lightning/requirements.in
+    #   -r benchmarks/llama/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   -r benchmarks/rlhf/requirements.in
+    #   -r benchmarks/timm/requirements.in
+    #   -r benchmarks/torchatari/requirements.in
+    #   -r benchmarks/torchvision/requirements.in
+    #   -r benchmarks/torchvision_ddp/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+    #   accelerate
+    #   botorch
+    #   diffusers
+    #   fairscale
+    #   lightning
+    #   linear-operator
+    #   pyro-ppl
+    #   pytorch-lightning
+    #   timm
+    #   torchmetrics
+    #   torchvision
+    #   xformers
+torch-cluster==1.6.3
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+torch-geometric==2.6.1
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+torch-scatter==2.1.2
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+torch-sparse==0.6.18
+    # via
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+torchao==0.8.0+rocm6.2.4
+    # via -r benchmarks/llm/requirements.in
+torchcompat==1.1.4
+    # via
+    #   -c /home/testroot/milabench/constraints/rocm.txt
+    #   -r benchmarks/flops/requirements.in
+    #   -r benchmarks/lightning/requirements.in
+    #   -r benchmarks/torchatari/requirements.in
+    #   -r benchmarks/torchvision/requirements.in
+    #   -r benchmarks/torchvision_ddp/requirements.in
+torchmetrics==1.6.1
+    # via
+    #   -r benchmarks/dinov2/requirements.in
+    #   lightning
+    #   pytorch-lightning
+torchtune==0.5.0
+    # via -r benchmarks/llm/requirements.in
+torchvision==0.21.0+rocm6.2.4
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/flops/requirements.in
+    #   -r benchmarks/lightning/requirements.in
+    #   -r benchmarks/timm/requirements.in
+    #   -r benchmarks/torchvision/requirements.in
+    #   -r benchmarks/torchvision_ddp/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+    #   timm
+tqdm==4.67.1
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/flops/requirements.in
+    #   -r benchmarks/torchvision/requirements.in
+    #   -r benchmarks/torchvision_ddp/requirements.in
+    #   datasets
+    #   fvcore
+    #   huggingface-hub
+    #   iopath
+    #   kagglehub
+    #   lightning
+    #   pyro-ppl
+    #   pytorch-lightning
+    #   torch-geometric
+    #   torchtune
+    #   transformers
+transformers==4.49.0
+    # via
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/huggingface/requirements.in
+    #   -r benchmarks/llama/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/rlhf/requirements.in
+    #   trl
+treescope==0.1.9
+    # via flax
+trimesh==4.6.3
+    # via
+    #   brax
+    #   mujoco-mjx
+trl==0.15.2
+    # via -r benchmarks/rlhf/requirements.in
+typeguard==4.4.2
+    # via tyro
+types-protobuf==5.29.1.20250208
+    # via envpool
+typing-extensions==4.12.2
+    # via
+    #   astroid
+    #   black
+    #   botorch
+    #   brax
+    #   chex
+    #   envpool
+    #   etils
+    #   flashbax
+    #   flax
+    #   gymnasium
+    #   huggingface-hub
+    #   iopath
+    #   lightning
+    #   lightning-utilities
+    #   multidict
+    #   navix
+    #   optree
+    #   orbax-checkpoint
+    #   pydantic
+    #   pydantic-core
+    #   pyre-extensions
+    #   pytorch-lightning
+    #   reactivex
+    #   rich
+    #   submitit
+    #   tables
+    #   torch
+    #   typeguard
+    #   typing-inspect
+    #   tyro
+    #   wandb
+typing-inspect==0.9.0
+    # via pyre-extensions
+tyro==0.9.16
+    # via
+    #   -r benchmarks/torchatari/requirements.in
+    #   navix
+tzdata==2025.1
+    # via pandas
+urllib3==2.3.0
+    # via
+    #   blobfile
+    #   requests
+    #   sentry-sdk
+varname==0.14.0
+    # via giving
+voir==0.2.19
+    # via
+    #   -c /home/testroot/milabench/constraints/rocm.txt
+    #   -r benchmarks/brax/requirements.in
+    #   -r benchmarks/diffusion/requirements.in
+    #   -r benchmarks/dinov2/requirements.in
+    #   -r benchmarks/flops/requirements.in
+    #   -r benchmarks/geo_gnn/requirements.in
+    #   -r benchmarks/huggingface/requirements.in
+    #   -r benchmarks/lightning/requirements.in
+    #   -r benchmarks/llama/requirements.in
+    #   -r benchmarks/llava/requirements.in
+    #   -r benchmarks/llm/requirements.in
+    #   -r benchmarks/purejaxrl/requirements.in
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   -r benchmarks/rlhf/requirements.in
+    #   -r benchmarks/timm/requirements.in
+    #   -r benchmarks/torchatari/requirements.in
+    #   -r benchmarks/torchvision/requirements.in
+    #   -r benchmarks/torchvision_ddp/requirements.in
+    #   -r benchmarks/vjepa/requirements.in
+wadler-lindig==0.1.3
+    # via jaxtyping
+wandb==0.19.7
+    # via
+    #   -r benchmarks/recursiongfn/requirements.in
+    #   navix
+webdataset==0.2.111
+    # via -r benchmarks/vjepa/requirements.in
+werkzeug==3.1.3
+    # via
+    #   flask
+    #   flask-cors
+    #   tensorboard
+wrapt==1.17.2
+    # via dm-tree
+xformers==0.0.29.post3
+    # via -r benchmarks/dinov2/requirements.in
+xxhash==3.5.0
+    # via datasets
+yacs==0.1.8
+    # via fvcore
+yarl==1.18.3
+    # via aiohttp
+zipp==3.21.0
+    # via
+    #   etils
+    #   importlib-metadata
+
+# The following packages are considered to be unsafe in a requirements file:
+# setuptools
diff --git a/benchmarks/brax/requirements.rocm.txt b/benchmarks/brax/requirements.rocm.txt
index e0a9eed6..aed701cf 100644
--- a/benchmarks/brax/requirements.rocm.txt
+++ b/benchmarks/brax/requirements.rocm.txt
@@ -4,16 +4,15 @@
 #
 #    pip-compile --output-file=benchmarks/brax/requirements.rocm.txt .pin/tmp-constraints-rocm-brax.txt benchmarks/brax/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 absl-py==2.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   dm-env
+    #   dm-tree
     #   ml-collections
     #   mujoco
     #   mujoco-mjx
@@ -21,115 +20,118 @@ absl-py==2.1.0
     #   orbax-checkpoint
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
+attrs==25.1.0
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
 blinker==1.9.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
 brax==0.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/brax/requirements.in
-chex==0.1.88
+chex==0.1.89
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   optax
 click==8.1.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
 cloudpickle==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
 dm-env==1.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-dm-tree==0.1.8
+dm-tree==0.1.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   dm-env
-etils[epath,epy]==1.11.0
+etils[epath,epy]==1.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco
     #   mujoco-mjx
     #   optax
     #   orbax-checkpoint
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 flask==3.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flask-cors
-flask-cors==5.0.0
+flask-cors==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-flax==0.10.2
+flax==0.10.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   etils
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 glfw==2.8.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   mujoco
-grpcio==1.69.0
+grpcio==1.70.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 gym==0.26.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 gym-notices==0.0.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
-humanize==4.11.0
+humanize==4.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   etils
 itsdangerous==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
-jax==0.4.38
+jax==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/brax/requirements.in
     #   brax
     #   chex
@@ -138,9 +140,9 @@ jax==0.4.38
     #   mujoco-mjx
     #   optax
     #   orbax-checkpoint
-jaxlib==0.4.38
+jaxlib==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   jax
@@ -149,69 +151,70 @@ jaxlib==0.4.38
     #   optax
 jaxopt==0.8.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flask
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
     #   werkzeug
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 ml-collections==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 ml-dtypes==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jax
     #   jaxlib
     #   tensorstore
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 msgpack==1.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   orbax-checkpoint
 mujoco==3.2.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco-mjx
 mujoco-mjx==3.2.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 nest-asyncio==1.6.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   dm-env
+    #   dm-tree
     #   gym
     #   jax
     #   jaxlib
@@ -223,131 +226,136 @@ numpy==1.26.4
     #   scipy
     #   tensorboardx
     #   tensorstore
+    #   treescope
     #   trimesh
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 opt-einsum==3.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jax
 optax==0.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flax
-orbax-checkpoint==0.11.0
+orbax-checkpoint==0.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flax
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboardx
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 protobuf==5.29.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
     #   tensorboardx
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pyopengl==3.1.7
+pyopengl==3.1.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   mujoco
 pytinyrenderer==0.0.14
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   ml-collections
     #   omegaconf
     #   orbax-checkpoint
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   voir
-scipy==1.15.1
+scipy==1.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   jax
     #   jaxlib
     #   jaxopt
     #   mujoco-mjx
-simplejson==3.19.3
+simplejson==3.20.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ml-collections
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tensorboardx==2.6.2.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-tensorstore==0.1.71
+tensorstore==0.1.72
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   orbax-checkpoint
 toolz==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   chex
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/brax/requirements.in
-trimesh==4.5.3
+treescope==0.1.9
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   flax
+trimesh==4.6.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco-mjx
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   etils
@@ -358,18 +366,23 @@ typing-extensions==4.12.2
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/brax/requirements.in
 werkzeug==3.1.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
+    #   flask-cors
+wrapt==1.17.2
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
 zipp==3.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   etils
diff --git a/benchmarks/diffusion/requirements.rocm.txt b/benchmarks/diffusion/requirements.rocm.txt
index 4a68dac9..f73066f6 100644
--- a/benchmarks/diffusion/requirements.rocm.txt
+++ b/benchmarks/diffusion/requirements.rocm.txt
@@ -4,109 +4,106 @@
 #
 #    pip-compile --output-file=benchmarks/diffusion/requirements.rocm.txt .pin/tmp-constraints-rocm-diffusion-nodes.txt benchmarks/diffusion/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-accelerate==1.2.1
+accelerate==1.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
     #   diffusers
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 argklass==1.4.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-datasets==3.2.0
+datasets==3.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
 diffusers[torch]==0.32.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   multiprocess
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   diffusers
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 hjson==3.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   diffusers
@@ -114,124 +111,124 @@ huggingface-hub==0.27.1
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
-importlib-metadata==8.5.0
+importlib-metadata==8.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   diffusers
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multiprocess==0.70.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   diffusers
     #   pandas
     #   torchvision
     #   transformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
     #   transformers
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   diffusers
     #   torchvision
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
@@ -239,97 +236,97 @@ pyyaml==6.0.2
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   diffusers
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   diffusers
     #   huggingface-hub
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   diffusers
     #   transformers
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   diffusers
     #   torchvision
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
     #   datasets
     #   huggingface-hub
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/diffusion/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   multidict
     #   reactivex
     #   rich
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/diffusion/requirements.in
 xxhash==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 zipp==3.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   importlib-metadata
diff --git a/benchmarks/dinov2/requirements.rocm.txt b/benchmarks/dinov2/requirements.rocm.txt
index ece66eaf..abf1fb06 100644
--- a/benchmarks/dinov2/requirements.rocm.txt
+++ b/benchmarks/dinov2/requirements.rocm.txt
@@ -4,189 +4,186 @@
 #
 #    pip-compile --output-file=benchmarks/dinov2/requirements.rocm.txt .pin/tmp-constraints-rocm-dinov2-giant-gpus.txt benchmarks/dinov2/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 cloudpickle==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   submitit
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 fvcore==0.1.5.post20221221
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 iopath==0.1.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
     #   fvcore
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-lightning-utilities==0.11.9
+lightning-utilities==0.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchmetrics
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
     #   scipy
     #   torchmetrics
     #   torchvision
     #   xformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning-utilities
     #   torchmetrics
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
     #   torchvision
 portalocker==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   iopath
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
     #   omegaconf
     #   yacs
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-scipy==1.15.1
+scipy==1.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
 submitit==1.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tabulate==0.9.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
 termcolor==2.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
     #   torchmetrics
     #   torchvision
     #   xformers
 torchmetrics==1.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
     #   iopath
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   iopath
     #   lightning-utilities
     #   reactivex
@@ -195,20 +192,20 @@ typing-extensions==4.12.2
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/dinov2/requirements.in
-xformers==0.0.29.post1
+xformers==0.0.29.post3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/dinov2/requirements.in
 yacs==0.1.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fvcore
 
 # The following packages are considered to be unsafe in a requirements file:
diff --git a/benchmarks/flops/requirements.rocm.txt b/benchmarks/flops/requirements.rocm.txt
index 8b310c1b..8297d303 100644
--- a/benchmarks/flops/requirements.rocm.txt
+++ b/benchmarks/flops/requirements.rocm.txt
@@ -4,150 +4,147 @@
 #
 #    pip-compile --output-file=benchmarks/flops/requirements.rocm.txt .pin/tmp-constraints-rocm-flops.txt benchmarks/flops/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchcompat
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/flops/requirements.in
     #   torchvision
 torchcompat==1.1.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/flops/requirements.in
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/flops/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/flops/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   reactivex
     #   rich
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/flops/requirements.in
diff --git a/benchmarks/geo_gnn/requirements-pre.rocm.txt b/benchmarks/geo_gnn/requirements-pre.rocm.txt
index 32085fd4..f2e206de 100644
--- a/benchmarks/geo_gnn/requirements-pre.rocm.txt
+++ b/benchmarks/geo_gnn/requirements-pre.rocm.txt
@@ -4,48 +4,45 @@
 #
 #    pip-compile --output-file=benchmarks/geo_gnn/requirements-pre.rocm.txt .pin/tmp-constraints-rocm-dimenet.txt benchmarks/geo_gnn/requirements-pre.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
diff --git a/benchmarks/geo_gnn/requirements.rocm.txt b/benchmarks/geo_gnn/requirements.rocm.txt
index 2a33b7f9..41b911dd 100644
--- a/benchmarks/geo_gnn/requirements.rocm.txt
+++ b/benchmarks/geo_gnn/requirements.rocm.txt
@@ -4,262 +4,259 @@
 #
 #    pip-compile --output-file=benchmarks/geo_gnn/requirements.rocm.txt .pin/tmp-constraints-rocm-dimenet.txt benchmarks/geo_gnn/requirements-pre.rocm.txt benchmarks/geo_gnn/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
-    #   pytorch-triton-rocm
     #   torch
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   torch
     #   torch-geometric
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   torch
     #   torch-geometric
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
     #   pandas
     #   rdkit
     #   scipy
     #   torch-geometric
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rdkit
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 pyparsing==3.2.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
-rdkit==2024.9.4
+rdkit==2024.9.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-scipy==1.15.1
+scipy==1.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-cluster
     #   torch-sparse
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
 torch-cluster==1.6.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 torch-geometric==2.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 torch-scatter==2.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 torch-sparse==0.6.18
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/geo_gnn/requirements-pre.rocm.txt
     #   multidict
     #   reactivex
     #   rich
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/geo_gnn/requirements.in
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
diff --git a/benchmarks/huggingface/requirements.rocm.txt b/benchmarks/huggingface/requirements.rocm.txt
index 1df97f95..684d3e4e 100644
--- a/benchmarks/huggingface/requirements.rocm.txt
+++ b/benchmarks/huggingface/requirements.rocm.txt
@@ -4,190 +4,187 @@
 #
 #    pip-compile --output-file=benchmarks/huggingface/requirements.rocm.txt .pin/tmp-constraints-rocm-hf.txt benchmarks/huggingface/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tokenizers
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   transformers
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/huggingface/requirements.in
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   omegaconf
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/huggingface/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/huggingface/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   reactivex
     #   rich
     #   torch
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/huggingface/requirements.in
diff --git a/benchmarks/lightning/requirements.rocm.txt b/benchmarks/lightning/requirements.rocm.txt
index b8e49c55..80b9558c 100644
--- a/benchmarks/lightning/requirements.rocm.txt
+++ b/benchmarks/lightning/requirements.rocm.txt
@@ -4,188 +4,185 @@
 #
 #    pip-compile --output-file=benchmarks/lightning/requirements.rocm.txt .pin/tmp-constraints-rocm-lightning-gpus.txt benchmarks/lightning/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   pytorch-lightning
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   yarl
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchcompat
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 lightning==2.5.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/lightning/requirements.in
-lightning-utilities==0.11.9
+lightning-utilities==0.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   pytorch-lightning
     #   torchmetrics
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchmetrics
     #   torchvision
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   lightning-utilities
     #   pytorch-lightning
     #   torchmetrics
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 pytorch-lightning==2.5.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   omegaconf
     #   pytorch-lightning
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/lightning/requirements.in
     #   lightning
     #   pytorch-lightning
@@ -193,26 +190,26 @@ torch==2.5.1+rocm6.2
     #   torchvision
 torchcompat==1.1.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/lightning/requirements.in
 torchmetrics==1.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   pytorch-lightning
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/lightning/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   pytorch-lightning
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   lightning
     #   lightning-utilities
     #   multidict
@@ -222,16 +219,16 @@ typing-extensions==4.12.2
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/lightning/requirements.in
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 
 # The following packages are considered to be unsafe in a requirements file:
diff --git a/benchmarks/llama/requirements.rocm.txt b/benchmarks/llama/requirements.rocm.txt
index 36d1b558..f43b4da1 100644
--- a/benchmarks/llama/requirements.rocm.txt
+++ b/benchmarks/llama/requirements.rocm.txt
@@ -4,296 +4,293 @@
 #
 #    pip-compile --output-file=benchmarks/llama/requirements.rocm.txt .pin/tmp-constraints-rocm-llm.txt benchmarks/llama/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-datasets==3.2.0
+datasets==3.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   multiprocess
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
 fairscale==0.4.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
 fire==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   tokenizers
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multiprocess==0.70.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fairscale
     #   pandas
     #   transformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   omegaconf
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 sentencepiece==0.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 termcolor==2.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fire
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
     #   fairscale
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llama/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   multidict
     #   reactivex
     #   rich
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/llama/requirements.in
 xxhash==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
diff --git a/benchmarks/llava/requirements.rocm.txt b/benchmarks/llava/requirements.rocm.txt
index 54515f81..e6d94665 100644
--- a/benchmarks/llava/requirements.rocm.txt
+++ b/benchmarks/llava/requirements.rocm.txt
@@ -4,209 +4,206 @@
 #
 #    pip-compile --output-file=benchmarks/llava/requirements.rocm.txt .pin/tmp-constraints-rocm-llava-single.txt benchmarks/llava/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-accelerate==1.2.1
+accelerate==1.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-datasets==3.2.0
+datasets==3.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   multiprocess
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   tokenizers
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multiprocess==0.70.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
     #   accelerate
     #   datasets
     #   pandas
     #   transformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
     #   transformers
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
@@ -214,84 +211,84 @@ pyyaml==6.0.2
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   transformers
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
     #   accelerate
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llava/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   multidict
     #   reactivex
     #   rich
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/llava/requirements.in
 xxhash==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
diff --git a/benchmarks/llm/requirements.rocm.txt b/benchmarks/llm/requirements.rocm.txt
index 59e657c8..29c3aab3 100644
--- a/benchmarks/llm/requirements.rocm.txt
+++ b/benchmarks/llm/requirements.rocm.txt
@@ -4,118 +4,119 @@
 #
 #    pip-compile --output-file=benchmarks/llm/requirements.rocm.txt .pin/tmp-constraints-rocm-llm-full-mp-nodes.txt benchmarks/llm/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-accelerate==1.2.1
+accelerate==1.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 argklass==1.4.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 blobfile==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -r benchmarks/llm/requirements.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
     #   torchtune
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-datasets==3.2.0
+datasets==3.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   multiprocess
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
 fairscale==0.4.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
     #   -r benchmarks/llm/requirements.in
-    #   -r benchmarks/llm/requirements.txt
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blobfile
     #   datasets
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
 fire==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -r benchmarks/llm/requirements.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
+hf-transfer==0.1.9
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   huggingface-hub
 hjson==3.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
-huggingface-hub[hf-transfer,hf_transfer]==0.27.1
+huggingface-hub[hf-transfer,hf_transfer]==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   tokenizers
@@ -123,79 +124,79 @@ huggingface-hub[hf-transfer,hf_transfer]==0.27.1
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-kagglehub==0.3.6
+kagglehub==0.3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
-lxml==5.3.0
+lxml==5.3.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blobfile
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multiprocess==0.70.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   fairscale
     #   pandas
     #   torchtune
     #   transformers
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
@@ -203,72 +204,73 @@ packaging==24.2
     #   transformers
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   torchtune
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pycryptodomex==3.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blobfile
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
     #   accelerate
     #   datasets
     #   huggingface-hub
+    #   kagglehub
     #   omegaconf
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tiktoken
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   kagglehub
@@ -276,96 +278,96 @@ requests==2.32.3
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   torchtune
     #   transformers
 sentencepiece==0.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 termcolor==2.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   fire
-tiktoken==0.8.0
+tiktoken==0.9.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchtune
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -r /home/testroot/milabench/benchmarks/llm/requirements.txt
     #   -r benchmarks/llm/requirements.in
-    #   -r benchmarks/llm/requirements.txt
     #   accelerate
     #   fairscale
-torchao==0.8.0
+torchao==0.8.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
 torchtune==0.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   kagglehub
     #   torchtune
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/llm/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   multidict
     #   reactivex
     #   rich
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blobfile
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/llm/requirements.in
 xxhash==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
diff --git a/benchmarks/purejaxrl/requirements.rocm.txt b/benchmarks/purejaxrl/requirements.rocm.txt
index b75c285d..dcd78c0e 100644
--- a/benchmarks/purejaxrl/requirements.rocm.txt
+++ b/benchmarks/purejaxrl/requirements.rocm.txt
@@ -4,17 +4,16 @@
 #
 #    pip-compile --output-file=benchmarks/purejaxrl/requirements.rocm.txt .pin/tmp-constraints-rocm-ppo.txt benchmarks/purejaxrl/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 absl-py==2.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   distrax
     #   dm-env
+    #   dm-tree
     #   ml-collections
     #   mujoco
     #   mujoco-mjx
@@ -24,48 +23,52 @@ absl-py==2.1.0
     #   tensorflow-probability
 annotated-types==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pydantic
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 argklass==1.4.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
 astroid==3.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pylint
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-black==24.10.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
+black==25.1.0
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
 blinker==1.9.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
 brax==0.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   sentry-sdk
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
-chex==0.1.88
+chex==0.1.89
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   distrax
     #   evosax
     #   flashbax
@@ -74,66 +77,66 @@ chex==0.1.88
     #   rlax
 click==8.1.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
     #   flask
     #   wandb
 cloudpickle==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
     #   gymnasium
     #   tensorflow-probability
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
 contourpy==1.3.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
 cycler==0.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
-decorator==5.1.1
+decorator==5.2.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorflow-probability
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pylint
 distrax==0.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
     #   rlax
 dm-env==1.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   rlax
-dm-tree==0.1.8
+dm-tree==0.1.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   dm-env
     #   tensorflow-probability
 docker-pycreds==0.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 docstring-parser==0.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
 dotmap==1.3.30
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   evosax
-etils[epath,epy]==1.11.0
+etils[epath,epy]==1.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco
     #   mujoco-mjx
@@ -141,135 +144,134 @@ etils[epath,epy]==1.11.0
     #   orbax-checkpoint
 evosax==0.1.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
 exceptiongroup==1.2.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pytest
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
 farama-notifications==0.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gymnasium
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-flake8==7.1.1
+flake8==7.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
 flashbax==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
 flask==3.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flask-cors
-flask-cors==5.0.0
+flask-cors==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-flax==0.10.2
+flax==0.10.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
     #   brax
     #   evosax
     #   flashbax
     #   gymnax
     #   navix
-fonttools==4.55.3
+fonttools==4.56.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   etils
     #   torch
 gast==0.6.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorflow-probability
 gitdb==4.0.12
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gitpython
 gitpython==3.1.44
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 glfw==2.8.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   mujoco
-grpcio==1.69.0
+grpcio==1.70.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 gym==0.26.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   gymnax
 gym-notices==0.0.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
-gymnasium==1.0.0
+gymnasium==1.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gymnax
 gymnax==0.0.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/purejaxrl/requirements.in
 hjson==3.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
-humanize==4.11.0
+humanize==4.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   argklass
     #   etils
 iniconfig==2.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pytest
-isort==5.13.2
+isort==6.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pylint
 itsdangerous==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
-jax==0.4.38
+jax==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
     #   brax
     #   chex
@@ -283,9 +285,9 @@ jax==0.4.38
     #   optax
     #   orbax-checkpoint
     #   rlax
-jaxlib==0.4.38
+jaxlib==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   chex
     #   distrax
@@ -299,95 +301,96 @@ jaxlib==0.4.38
     #   rlax
 jaxopt==0.8.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flask
     #   torch
 kiwisolver==1.4.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
     #   werkzeug
 matplotlib==3.10.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   evosax
     #   gymnax
     #   seaborn
 mccabe==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flake8
     #   pylint
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 ml-collections==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 ml-dtypes==0.5.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jax
     #   jaxlib
     #   tensorstore
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 msgpack==1.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   orbax-checkpoint
 mujoco==3.2.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco-mjx
 mujoco-mjx==3.2.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 mypy-extensions==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
 navix==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
 nest-asyncio==1.6.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
     #   brax
     #   chex
     #   contourpy
     #   distrax
     #   dm-env
+    #   dm-tree
     #   evosax
     #   flashbax
     #   gym
@@ -408,37 +411,38 @@ numpy==1.26.4
     #   tensorboardx
     #   tensorflow-probability
     #   tensorstore
+    #   treescope
     #   trimesh
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 opt-einsum==3.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jax
 optax==0.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
     #   brax
     #   flax
-orbax-checkpoint==0.11.0
+orbax-checkpoint==0.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   flax
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
     #   matplotlib
     #   pytest
@@ -446,99 +450,99 @@ packaging==24.2
     #   tensorboardx
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   seaborn
 pathspec==0.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   matplotlib
     #   navix
 platformdirs==4.3.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
     #   pylint
     #   wandb
 pluggy==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pytest
 protobuf==5.29.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
     #   tensorboardx
     #   wandb
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
     #   wandb
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pycodestyle==2.12.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flake8
-pydantic==2.10.5
+pydantic==2.10.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 pydantic-core==2.27.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pydantic
 pyflakes==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flake8
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pylint==3.3.3
+pylint==3.3.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
-pyopengl==3.1.7
+pyopengl==3.1.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   mujoco
 pyparsing==3.2.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
 pytest==8.3.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   matplotlib
     #   pandas
 pytinyrenderer==0.0.14
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   evosax
     #   flax
     #   gymnax
@@ -548,25 +552,25 @@ pyyaml==6.0.2
     #   wandb
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flax
     #   tyro
     #   voir
 rlax==0.1.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
-scipy==1.15.1
+scipy==1.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   jax
     #   jaxlib
@@ -574,88 +578,92 @@ scipy==1.15.1
     #   mujoco-mjx
 seaborn==0.13.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gymnax
-sentry-sdk==2.20.0
+sentry-sdk==2.22.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
-setproctitle==1.3.4
+setproctitle==1.3.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
-setuptools-scm==8.1.0
+setuptools-scm==8.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
 shtab==1.7.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
-simplejson==3.19.3
+simplejson==3.20.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   orbax-checkpoint
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   docker-pycreds
     #   ml-collections
     #   python-dateutil
     #   tensorflow-probability
 smmap==5.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gitdb
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tensorboardx==2.6.2.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
 tensorflow-probability==0.25.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   distrax
-tensorstore==0.1.71
+tensorstore==0.1.72
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flashbax
     #   flax
     #   orbax-checkpoint
 tomli==2.2.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   black
     #   pylint
     #   pytest
     #   setuptools-scm
 tomlkit==0.13.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pylint
 toolz==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   chex
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/purejaxrl/requirements.in
-trimesh==4.5.3
+treescope==0.1.9
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   flax
+trimesh==4.6.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   brax
     #   mujoco-mjx
-typeguard==4.4.1
+typeguard==4.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   astroid
     #   black
     #   brax
@@ -674,39 +682,44 @@ typing-extensions==4.12.2
     #   typeguard
     #   tyro
     #   wandb
-tyro==0.9.10
+tyro==0.9.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   sentry-sdk
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/purejaxrl/requirements.in
-wandb==0.19.4
+wandb==0.19.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   navix
 werkzeug==3.1.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   flask
+    #   flask-cors
+wrapt==1.17.2
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
 zipp==3.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   etils
 
 # The following packages are considered to be unsafe in a requirements file:
diff --git a/benchmarks/recursiongfn/requirements.rocm.txt b/benchmarks/recursiongfn/requirements.rocm.txt
index 3f8f2f44..ec19d7cb 100644
--- a/benchmarks/recursiongfn/requirements.rocm.txt
+++ b/benchmarks/recursiongfn/requirements.rocm.txt
@@ -4,201 +4,200 @@
 #
 #    pip-compile --output-file=benchmarks/recursiongfn/requirements.rocm.txt .pin/tmp-constraints-rocm-recursiongfn.txt benchmarks/recursiongfn/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 absl-py==2.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 annotated-types==0.7.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pydantic
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 blosc2==2.7.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tables
-botorch==0.12.0
+botorch==0.13.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   sentry-sdk
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 click==8.1.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
 cvxopt==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 docker-pycreds==0.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
     #   torch-geometric
 gitdb==4.0.12
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gitpython
 gitpython==3.1.44
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   wandb
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-gpytorch==1.13
+gpytorch==1.14
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   botorch
-grpcio==1.69.0
+grpcio==1.70.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
-jaxtyping==0.2.19
+jaxtyping==0.2.38
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gpytorch
     #   linear-operator
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
     #   torch-geometric
 joblib==1.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   scikit-learn
-linear-operator==0.5.3
+linear-operator==0.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   botorch
     #   gpytorch
 markdown==3.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
     #   werkzeug
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   botorch
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gpytorch
     #   linear-operator
     #   sympy
 msgpack==1.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blosc2
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multipledispatch==1.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   botorch
+mypy-extensions==1.0.0
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   typing-inspect
 ndindex==1.9.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blosc2
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   torch
 numexpr==2.10.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blosc2
     #   tables
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blosc2
-    #   jaxtyping
     #   numexpr
     #   pandas
     #   pyro-ppl
@@ -208,135 +207,139 @@ numpy==1.26.4
     #   tables
     #   tensorboard
     #   torch-geometric
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   voir
 opt-einsum==3.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pyro-ppl
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tables
     #   tensorboard
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rdkit
 platformdirs==4.3.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 protobuf==5.29.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
     #   wandb
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
     #   voir
     #   wandb
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 py-cpuinfo==9.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   blosc2
     #   tables
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
-pydantic==2.10.5
+pydantic==2.10.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 pydantic-core==2.27.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pydantic
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 pyparsing==3.2.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
+pyre-extensions==0.0.32
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   botorch
 pyro-api==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pyro-ppl
 pyro-ppl==1.9.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   botorch
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
     #   wandb
-rdkit==2024.9.4
+rdkit==2024.9.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch-geometric
     #   wandb
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 scikit-learn==1.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gpytorch
-scipy==1.15.1
+scipy==1.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   botorch
     #   gpytorch
@@ -344,119 +347,124 @@ scipy==1.15.1
     #   scikit-learn
     #   torch-cluster
     #   torch-sparse
-sentry-sdk==2.20.0
+sentry-sdk==2.22.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
-setproctitle==1.3.4
+setproctitle==1.3.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   wandb
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   docker-pycreds
     #   python-dateutil
     #   tensorboard
 smmap==5.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gitdb
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tables==3.10.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
-tensorboard==2.18.0
+tensorboard==2.19.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 tensorboard-data-server==0.7.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 threadpoolctl==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   botorch
     #   scikit-learn
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
     #   botorch
     #   linear-operator
     #   pyro-ppl
 torch-cluster==1.6.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 torch-geometric==2.6.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 torch-scatter==2.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 torch-sparse==0.6.18
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pyro-ppl
     #   torch-geometric
-typeguard==4.4.1
-    # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   jaxtyping
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   botorch
-    #   jaxtyping
     #   multidict
     #   pydantic
     #   pydantic-core
+    #   pyre-extensions
     #   reactivex
     #   rich
     #   tables
     #   torch
-    #   typeguard
+    #   typing-inspect
     #   wandb
-tzdata==2024.2
+typing-inspect==0.9.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   pyre-extensions
+tzdata==2025.1
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   sentry-sdk
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/recursiongfn/requirements.in
-wandb==0.19.4
+wadler-lindig==0.1.3
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   jaxtyping
+wandb==0.19.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/recursiongfn/requirements.in
 werkzeug==3.1.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 
 # The following packages are considered to be unsafe in a requirements file:
diff --git a/benchmarks/rlhf/requirements.rocm.txt b/benchmarks/rlhf/requirements.rocm.txt
index 6d61676e..9e91a184 100644
--- a/benchmarks/rlhf/requirements.rocm.txt
+++ b/benchmarks/rlhf/requirements.rocm.txt
@@ -4,215 +4,207 @@
 #
 #    pip-compile --output-file=benchmarks/rlhf/requirements.rocm.txt .pin/tmp-constraints-rocm-rlhf-gpus.txt benchmarks/rlhf/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
-accelerate==1.2.1
+accelerate==1.4.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
     #   trl
-aiohappyeyeballs==2.4.4
+aiohappyeyeballs==2.4.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-aiohttp==3.11.11
+aiohttp==3.11.13
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   fsspec
 aiosignal==1.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 async-timeout==5.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-attrs==24.3.0
+attrs==25.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-datasets==3.2.0
+datasets==3.3.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
     #   trl
 dill==0.3.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   multiprocess
-docstring-parser==0.16
+einops==0.8.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   tyro
-einops==0.8.0
-    # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
     #   transformers
 frozenlist==1.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   aiosignal
-fsspec[http]==2024.9.0
+fsspec[http]==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   tokenizers
     #   transformers
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
     #   yarl
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 multidict==6.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 multiprocess==0.70.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   pandas
     #   transformers
-    #   trl
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
     #   transformers
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
-propcache==0.2.1
+propcache==0.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
     #   yarl
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-pyarrow==19.0.0
+pyarrow==19.0.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   datasets
     #   huggingface-hub
@@ -220,106 +212,90 @@ pyyaml==6.0.2
     #   transformers
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 regex==2024.11.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   tyro
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   trl
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   accelerate
     #   transformers
-shtab==1.7.1
-    # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   tyro
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 tokenizers==0.21.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   transformers
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
     #   accelerate
-    #   trl
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
     #   huggingface-hub
     #   transformers
-transformers==4.48.0
+transformers==4.49.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
     #   trl
-trl==0.15.1
+trl==0.15.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/rlhf/requirements.in
-typeguard==4.4.1
-    # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   tyro
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   multidict
     #   reactivex
     #   rich
     #   torch
-    #   typeguard
-    #   tyro
-tyro==0.9.10
-    # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   trl
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/rlhf/requirements.in
 xxhash==3.5.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   datasets
 yarl==1.18.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   aiohttp
diff --git a/benchmarks/timm/requirements.rocm.txt b/benchmarks/timm/requirements.rocm.txt
index 4cd62365..c493a3f6 100644
--- a/benchmarks/timm/requirements.rocm.txt
+++ b/benchmarks/timm/requirements.rocm.txt
@@ -4,178 +4,175 @@
 #
 #    pip-compile --output-file=benchmarks/timm/requirements.rocm.txt .pin/tmp-constraints-rocm-timm.txt benchmarks/timm/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/timm/requirements.in
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/timm/requirements.in
     #   huggingface-hub
     #   omegaconf
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/timm/requirements.in
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/timm/requirements.in
     #   torchvision
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/timm/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   reactivex
     #   rich
     #   torch
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/timm/requirements.in
diff --git a/benchmarks/torchatari/requirements.rocm.txt b/benchmarks/torchatari/requirements.rocm.txt
index 8ec1a917..c01583b3 100644
--- a/benchmarks/torchatari/requirements.rocm.txt
+++ b/benchmarks/torchatari/requirements.rocm.txt
@@ -4,232 +4,235 @@
 #
 #    pip-compile --output-file=benchmarks/torchatari/requirements.rocm.txt .pin/tmp-constraints-rocm-torchatari.txt benchmarks/torchatari/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 absl-py==2.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   dm-env
+    #   dm-tree
     #   tensorboard
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 appdirs==1.4.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   cantilever
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
+attrs==25.1.0
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
 cantilever==0.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
 cloudpickle==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
     #   gymnasium
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
 dm-env==1.6
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
-dm-tree==0.1.8
+dm-tree==0.1.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   dm-env
 docstring-parser==0.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
 envpool==0.8.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
 farama-notifications==0.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gymnasium
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-grpcio==1.69.0
+grpcio==1.70.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 gym==0.26.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
     #   envpool
 gym-notices==0.0.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   gym
-gymnasium==1.0.0
+gymnasium==1.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   cantilever
     #   torchcompat
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown==3.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
     #   werkzeug
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
     #   dm-env
+    #   dm-tree
     #   envpool
     #   gym
     #   gymnasium
     #   tensorboard
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 optree==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
     #   tensorboard
 protobuf==5.29.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
     #   voir
 shtab==1.7.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-tensorboard==2.18.0
+tensorboard==2.19.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
 tensorboard-data-server==0.7.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
 torchcompat==1.1.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchatari/requirements.in
-typeguard==4.4.1
+typeguard==4.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tyro
-types-protobuf==5.29.1.20241207
+types-protobuf==5.29.1.20250208
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   envpool
     #   gymnasium
     #   optree
@@ -238,23 +241,27 @@ typing-extensions==4.12.2
     #   torch
     #   typeguard
     #   tyro
-tyro==0.9.10
+tyro==0.9.16
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchatari/requirements.in
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchatari/requirements.in
 werkzeug==3.1.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   tensorboard
+wrapt==1.17.2
+    # via
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   dm-tree
 
 # The following packages are considered to be unsafe in a requirements file:
 # setuptools
diff --git a/benchmarks/torchvision/requirements.rocm.txt b/benchmarks/torchvision/requirements.rocm.txt
index 7d746afb..f9d737fb 100644
--- a/benchmarks/torchvision/requirements.rocm.txt
+++ b/benchmarks/torchvision/requirements.rocm.txt
@@ -4,150 +4,147 @@
 #
 #    pip-compile --output-file=benchmarks/torchvision/requirements.rocm.txt .pin/tmp-constraints-rocm-torchvision.txt benchmarks/torchvision/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchcompat
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision/requirements.in
     #   torchvision
 torchcompat==1.1.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchvision/requirements.in
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   reactivex
     #   rich
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchvision/requirements.in
diff --git a/benchmarks/torchvision_ddp/requirements.rocm.txt b/benchmarks/torchvision_ddp/requirements.rocm.txt
index 5e54bb9d..22400c72 100644
--- a/benchmarks/torchvision_ddp/requirements.rocm.txt
+++ b/benchmarks/torchvision_ddp/requirements.rocm.txt
@@ -4,150 +4,147 @@
 #
 #    pip-compile --output-file=benchmarks/torchvision_ddp/requirements.rocm.txt .pin/tmp-constraints-rocm-torchvision.txt benchmarks/torchvision_ddp/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   pytorch-triton-rocm
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
 importlib-resources==6.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchcompat
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision_ddp/requirements.in
     #   torchvision
 torchcompat==1.1.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchvision_ddp/requirements.in
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision_ddp/requirements.in
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/torchvision_ddp/requirements.in
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   reactivex
     #   rich
     #   torch
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/torchvision_ddp/requirements.in
diff --git a/benchmarks/vjepa/requirements.rocm.txt b/benchmarks/vjepa/requirements.rocm.txt
index ee7d650d..0ffac47b 100644
--- a/benchmarks/vjepa/requirements.rocm.txt
+++ b/benchmarks/vjepa/requirements.rocm.txt
@@ -4,167 +4,164 @@
 #
 #    pip-compile --output-file=benchmarks/vjepa/requirements.rocm.txt .pin/tmp-constraints-rocm-vjepa-gpus.txt benchmarks/vjepa/requirements.in
 #
---extra-index-url https://pypi.ngc.nvidia.com
---extra-index-url https://download.pytorch.org/whl/rocm6.2
---trusted-host pypi.ngc.nvidia.com
+--extra-index-url https://download.pytorch.org/whl/rocm6.2.4
 
 antlr4-python3-runtime==4.9.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   omegaconf
 asttokens==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-beartype==0.19.0
+beartype==0.20.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
 braceexpand==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
     #   webdataset
-certifi==2024.12.14
+certifi==2025.1.31
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 charset-normalizer==3.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 cloudpickle==3.1.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   submitit
 codefind==0.1.7
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
 decord==0.6.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
-einops==0.8.0
+einops==0.8.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
-executing==2.1.0
+executing==2.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   varname
-filelock==3.16.1
+filelock==3.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
-    #   pytorch-triton-rocm
     #   torch
-fsspec==2024.9.0
+fsspec==2024.12.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   torch
 giving==0.4.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   ptera
     #   voir
-huggingface-hub==0.27.1
+huggingface-hub==0.29.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   timm
 idna==3.10
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 jinja2==3.1.5
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 markdown-it-py==3.0.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 markupsafe==3.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   jinja2
 mdurl==0.1.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   markdown-it-py
 mpmath==1.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   sympy
 networkx==3.4.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
 numpy==1.26.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
     #   decord
     #   opencv-python
     #   pandas
     #   torchvision
     #   webdataset
-nvidia-ml-py==12.560.30
+nvidia-ml-py==12.570.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 omegaconf==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 opencv-python==4.11.0.86
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
 ovld==0.3.9
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 packaging==24.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 pandas==2.2.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
 pillow==11.1.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torchvision
 psutil==5.9.8
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 ptera==1.4.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
 pygments==2.19.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   rich
 python-dateutil==2.9.0.post0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
-pytorch-triton-rocm==3.1.0
+pytorch-triton-rocm==3.2.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-pytz==2024.2
+pytz==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 pyyaml==6.0.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
     #   huggingface-hub
     #   omegaconf
@@ -172,77 +169,77 @@ pyyaml==6.0.2
     #   webdataset
 reactivex==4.0.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 requests==2.32.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 rich==13.9.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   voir
-safetensors==0.5.2
+safetensors==0.5.3
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   timm
 six==1.17.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   python-dateutil
 submitit==1.5.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
 sympy==1.13.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   torch
-timm==1.0.13
+timm==1.0.15
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
-torch==2.5.1+rocm6.2
+torch==2.6.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
     #   timm
     #   torchvision
-torchvision==0.20.1+rocm6.2
+torchvision==0.21.0+rocm6.2.4
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
     #   timm
 tqdm==4.67.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
 typing-extensions==4.12.2
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   huggingface-hub
     #   reactivex
     #   rich
     #   submitit
     #   torch
-tzdata==2024.2
+tzdata==2025.1
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   pandas
 urllib3==2.3.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   requests
 varname==0.14.0
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
 voir==0.2.19
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
-    #   -c .pin/../constraints/rocm.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/vjepa/requirements.in
-webdataset==0.2.100
+webdataset==0.2.111
     # via
-    #   -c .pin/../.pin/constraints-rocm-torch.txt
+    #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -r benchmarks/vjepa/requirements.in
diff --git a/benchmate/benchmate/warden.py b/benchmate/benchmate/warden.py
index a67681b0..a73a0287 100644
--- a/benchmate/benchmate/warden.py
+++ b/benchmate/benchmate/warden.py
@@ -75,8 +75,14 @@ def _default():
 def _rocm_parse_processes():
     cmd = ["rocm-smi", f"--showpids", "--json"]
     output = subprocess.check_output(cmd, text=True)
-    data = json.loads(output)
+    
     info = []
+
+    try:
+        data = json.loads(output)
+    except json.decoder.JSONDecodeError:
+        return info
+
     for key, data in data.get("system", {}).items():
         process_name, ngpu, vram, sdma, cu_occupancy = data.split(",")
         pid = key[3:]
diff --git a/milabench/_version.py b/milabench/_version.py
index d109cfd9..ae5133d5 100644
--- a/milabench/_version.py
+++ b/milabench/_version.py
@@ -1,5 +1,5 @@
 """This file is generated, do not modify"""
 
-__tag__ = "v1.0.0-50-g3249edf"
-__commit__ = "3249edfd466030ba2d7a612ac8b8102f6253a5a7"
-__date__ = "2025-02-06 10:51:24 -0500"
+__tag__ = "v1.0.0-54-g5440826"
+__commit__ = "5440826e4cdd300ce0a5c051d03bb6233342d954"
+__date__ = "2025-02-26 19:26:11 +0000"
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index 7551b97c..dbf73da9 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -124,8 +124,9 @@ install_prepare() {
     #
     # Install milabench's benchmarks in their venv
     #
-    # pip install torch --index-url https://download.pytorch.org/whl/rocm6.1
-    milabench pin --variant rocm --from-scratch $ARGS 
+    # pip install torch --index-url https://download.pytorch.org/whl/rocm6.2.4
+    # milabench pin --variant rocm --from-scratch $ARGS 
+    
     milabench install $ARGS 
 
     #
@@ -165,7 +166,7 @@ fi
 
 #
 #   Run the benchmakrs
-milabench run $ARGS # --system $MILABENCH_WORDIR/system.yaml
+# milabench run $ARGS # --system $MILABENCH_WORDIR/system.yaml
 
 
 #

From e76e0e4180eb180e1ad6a94055b3d1ee3a4692e7 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Thu, 27 Feb 2025 15:44:55 +0000
Subject: [PATCH 4/7] Add optimized lookup, add accelerator driver info, tweak
 llm prepare

---
 benchmarks/llm/prepare.py     |  4 +++-
 benchmate/benchmate/warden.py | 13 ++++++++++---
 config/base.yaml              | 11 +++++++----
 milabench/metadata.py         | 17 +++++++++++++++--
 milabench/sizer.py            | 32 +++++++++++++++++++++++++++++---
 scripts/article/run_rocm.sh   | 16 ++++++++++++++--
 6 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/benchmarks/llm/prepare.py b/benchmarks/llm/prepare.py
index 5f74d2a4..a5b633a0 100755
--- a/benchmarks/llm/prepare.py
+++ b/benchmarks/llm/prepare.py
@@ -150,7 +150,7 @@ def main():
 
     repo_id = config["repo_id"]
     hf_token = os.getenv("MILABENCH_HF_TOKEN", None)
-    output_dir = config["checkpointer"]["output_dir"]
+    output_dir = config["checkpointer"]["checkpoint_dir"]
 
     #
     huggingface_format = config.get("safetensors", False)
@@ -168,6 +168,8 @@ def main():
         # Ignore hugging face weights
         ignore_patterns = ["*.safetensors"]
 
+    print(f"Downloading to {output_dir}")
+
     download_args = [
         "download",
         repo_id,
diff --git a/benchmate/benchmate/warden.py b/benchmate/benchmate/warden.py
index a73a0287..7a1a9dbe 100644
--- a/benchmate/benchmate/warden.py
+++ b/benchmate/benchmate/warden.py
@@ -167,14 +167,21 @@ def children_warden(enabled=True):
         yield
         return
 
-    yield
-
-    pid = os.getpid()
 
     def get_children():
         with open(f"/proc/{pid}/task/{pid}/children", "r") as f:
             return [int(c) for c in f.read().strip().split()]
 
+    prev = set(get_children())
+
+    def get_children():
+        with open(f"/proc/{pid}/task/{pid}/children", "r") as f:
+            return set([int(c) for c in f.read().strip().split()]) - prev
+
+    yield
+
+    pid = os.getpid()
+
     def wait_for_children():
         children = get_children()
         start_time = time.time()
diff --git a/config/base.yaml b/config/base.yaml
index d612db23..a19de2a1 100644
--- a/config/base.yaml
+++ b/config/base.yaml
@@ -580,6 +580,7 @@ llm-lora-single:
     tokenizer.path={milabench_data}/llama3_8B/original/tokenizer.model: true
     safetensors=true: true
     checkpointer.checkpoint_dir={milabench_data}/llama3_8B: true
+    checkpointer.output_dir={milabench_extra}/llama3_8B: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
     batch_size=auto_batch(8): true
@@ -605,6 +606,7 @@ llm-lora-ddp-gpus:
     output_dir={milabench_extra}/output: true
     tokenizer.path={milabench_data}/llama3_8B/original/tokenizer.model: true
     checkpointer.checkpoint_dir={milabench_data}/llama3_8B/: true
+    checkpointer.output_dir={milabench_extra}/llama3_8B/: true
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
@@ -630,8 +632,8 @@ llm-lora-ddp-nodes:
     epochs=1: true
     output_dir={milabench_extra}/output: true
     tokenizer.path={milabench_data}/llama3_8B/original/tokenizer.model: true
-    checkpointer.checkpoint_dir={milabench_data}/llama3_8B/original: true
-    checkpointer.output_dir={milabench_data}/llama3_8B/: true
+    checkpointer.checkpoint_dir={milabench_data}/llama3_8B: true
+    checkpointer.output_dir={milabench_extra}/llama3_8B/: true
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-8B": true
@@ -659,6 +661,7 @@ llm-lora-mp-gpus:
     output_dir={milabench_extra}/output: true
     tokenizer.path={milabench_data}/llama3_70B/original/tokenizer.model: true
     checkpointer.checkpoint_dir={milabench_data}/llama3_70B: true
+    checkpointer.output_dir={milabench_extra}/llama3_70B/: true
     safetensors=true: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
@@ -687,7 +690,7 @@ llm-full-mp-gpus:
     output_dir={milabench_extra}/output: true
     tokenizer.path={milabench_data}/llama3_70B/original/tokenizer.model: true
     checkpointer.checkpoint_dir={milabench_data}/llama3_70B: true
-    checkpointer.output_dir={milabench_data}/llama3_70B_out/: true
+    checkpointer.output_dir={milabench_extra}/llama3_70B: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
     safetensors=true: true
@@ -712,7 +715,7 @@ llm-full-mp-nodes:
     output_dir={milabench_extra}/output: true
     tokenizer.path={milabench_data}/llama3_70B/original/tokenizer.model: true
     checkpointer.checkpoint_dir={milabench_data}/llama3_70B: true
-    checkpointer.output_dir={milabench_data}/llama3_70B/: true
+    checkpointer.output_dir={milabench_extra}/llama3_70B/: true
     metric_logger.log_dir={milabench_extra}/metrics: true
     repo_id="meta-llama/Meta-Llama-3.1-70B": true
     safetensors=true: true
diff --git a/milabench/metadata.py b/milabench/metadata.py
index eba6891b..92e92258 100644
--- a/milabench/metadata.py
+++ b/milabench/metadata.py
@@ -6,7 +6,7 @@
 import functools
 
 import cpuinfo
-from voir.instruments.gpu import get_gpu_info
+import voir.instruments.gpu as gpu
 
 import milabench.scripts.torchversion as torchversion
 
@@ -14,9 +14,16 @@
 from .scripts.vcs import retrieve_git_versions
 from .utils import error_guard
 
+
 def _get_gpu_info():
     try:
-        return get_gpu_info()
+        smi = gpu.select_backend()
+
+        return {
+            "arch": smi.arch,
+            "gpus": smi.get_gpus_info(),
+            "system": smi.system_info(),
+        }
     except Exception:
         traceback.print_exc()
         return {}
@@ -72,3 +79,9 @@ def machine_metadata(pack=None):
         ),
         "pytorch": torchv,
     }
+
+
+if __name__ == "__main__":
+    import json
+
+    print(json.dumps(machine_metadata(), indent=2))
diff --git a/milabench/sizer.py b/milabench/sizer.py
index 0bec77bb..a64e6a9a 100644
--- a/milabench/sizer.py
+++ b/milabench/sizer.py
@@ -133,8 +133,9 @@ def _scaling_v2(self, config):
 
         mem = [to_octet(v["memory"]) for v in data]
         size = [float(v["batch_size"]) for v in data]
+        perf = [float(v["perf"]) for v in data]
 
-        return mem, size
+        return mem, size, perf
 
     def get_capacity(self, capacity):
         if self.options.capacity is not None:
@@ -165,7 +166,7 @@ def auto_size(self, benchmark, capacity):
         if "model" in config:
             mem, size = self._scaling_v1(config)
         else:
-            mem, size = self._scaling_v2(config)
+            mem, size, _ = self._scaling_v2(config)
 
         if len(mem) == 1:
             syslog(f"Not enough data for {benchmark.config['name']}")
@@ -195,6 +196,31 @@ def auto_size(self, benchmark, capacity):
 
         return max(final_size, 1)
 
+    def optimized(self, benchmark, capacity):
+        # Old V1 format
+        config = self.benchscaling(benchmark)
+        if "model" in config:
+            return config["optimized"]
+
+        # Look for the best batch size
+        capacity = self.get_capacity(capacity)
+
+        if capacity is None:
+            syslog("Capacity is missing")
+            return None
+
+        # Look for the best batch size, make sure the used memory can fit
+        # in our current GPU
+        data = config["observations"]
+        data = list(sorted(data, key=lambda x: x["perf"], reverse=True))
+
+        for obs in data:
+            used_mem = to_octet(v["memory"])
+            if used_men < capacity:
+                return int(obs["batch_size"])
+        
+        return None
+
     def size(self, benchmark, capacity):
         config = self.benchscaling(benchmark)
 
@@ -202,7 +228,7 @@ def size(self, benchmark, capacity):
             return self.options.size
 
         if self.options.optimized:
-            return config["optimized"]
+            return self.optimized(benchmark, capacity)
 
         if self.options.autoscale:
             return self.auto_size(benchmark, capacity)
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index dbf73da9..916c2a22 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -118,6 +118,7 @@ install_prepare() {
     fi
 
     . $MILABENCH_WORDIR/env/bin/activate
+
     pip install -e $MILABENCH_SOURCE
 
     
@@ -154,19 +155,30 @@ fi
 
 
 
+(
+    . $MILABENCH_WORDIR/env/bin/activate
+    pip install git+https://github.com/breuleux/voir.git
+
+    (
+        . $BENCHMARK_VENV/bin/activate
+        pip install git+https://github.com/breuleux/voir.git
+    )
+
+)
+
 # (
 #     # . $BENCHMARK_VENV/bin/activate
 #     # pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
 # )
 
-# milabench install $ARGS --system $MILABENCH_WORDIR/system.yaml
+milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
 
 # milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
 # install_prepare
 
 #
 #   Run the benchmakrs
-# milabench run $ARGS # --system $MILABENCH_WORDIR/system.yaml
+milabench run $ARGS # --system $MILABENCH_WORDIR/system.yaml
 
 
 #

From 38c0015543c78622b555c2aad540bb499bcd8c12 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Thu, 27 Feb 2025 17:13:31 +0000
Subject: [PATCH 5/7] Update optimized lookup

---
 benchmarks/geo_gnn/pcqm4m_subset.py         |  2 +-
 benchmarks/llm/configs/llama3_70B_full.yaml |  2 +-
 milabench/sizer.py                          | 12 +++++++++---
 scripts/article/run_rocm.sh                 | 18 +++++++++---------
 4 files changed, 20 insertions(+), 14 deletions(-)

diff --git a/benchmarks/geo_gnn/pcqm4m_subset.py b/benchmarks/geo_gnn/pcqm4m_subset.py
index cee88c43..eb2c6f43 100644
--- a/benchmarks/geo_gnn/pcqm4m_subset.py
+++ b/benchmarks/geo_gnn/pcqm4m_subset.py
@@ -41,7 +41,7 @@ def __init__(
         self.from_smiles = from_smiles or _from_smiles
         super(PCQM4Mv2, self).__init__(root, transform, backend=backend, schema=schema)
 
-        split_idx = torch.load(self.raw_paths[1])
+        split_idx = torch.load(self.raw_paths[1], weights_only=False)
         self._indices = split_idx[self.split_mapping[split]].tolist()
 
     def raw_file_names(self):
diff --git a/benchmarks/llm/configs/llama3_70B_full.yaml b/benchmarks/llm/configs/llama3_70B_full.yaml
index 0032cf9d..1119d6a3 100644
--- a/benchmarks/llm/configs/llama3_70B_full.yaml
+++ b/benchmarks/llm/configs/llama3_70B_full.yaml
@@ -72,7 +72,7 @@ enable_activation_checkpointing: True  # True reduces memory
 enable_activation_offloading: True  # True reduces memory
 # custom_sharded_layers: ['tok_embeddings', 'output']  # Layers to shard separately (useful for large vocab size models). Lower Memory, but lower speed.
 # Hangs on AMD
-fsdp_cpu_offload: False
+# fsdp_cpu_offload: True
 compile: False  # torch.compile the model + loss, True increases speed + decreases memory
 optimizer_in_bwd: False  # True saves memory. Requires gradient_accumulation_steps=1
 
diff --git a/milabench/sizer.py b/milabench/sizer.py
index a64e6a9a..ac8fe03d 100644
--- a/milabench/sizer.py
+++ b/milabench/sizer.py
@@ -22,12 +22,17 @@
 default_scaling_folder = os.path.join(ROOT, "..", "config", "scaling")
 default_scaling_config = os.path.join(default_scaling_folder, "default.yaml")
 
+gpu_name_to_file = {
+    "AMD Instinct MI325 OAM": "MI325",
+}
+
 
 def gpu_name():
     try:
         info = get_gpu_info()
         values = list(info["gpus"].values())
-        return values[0]["product"]
+        n = values[0]["product"]
+        return gpu_name_to_file.get(n, n)
     except:
         return None
 
@@ -199,6 +204,7 @@ def auto_size(self, benchmark, capacity):
     def optimized(self, benchmark, capacity):
         # Old V1 format
         config = self.benchscaling(benchmark)
+
         if "model" in config:
             return config["optimized"]
 
@@ -215,8 +221,8 @@ def optimized(self, benchmark, capacity):
         data = list(sorted(data, key=lambda x: x["perf"], reverse=True))
 
         for obs in data:
-            used_mem = to_octet(v["memory"])
-            if used_men < capacity:
+            used_mem = to_octet(obs["memory"])
+            if used_mem < capacity:
                 return int(obs["batch_size"])
         
         return None
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index 916c2a22..85f1c8b3 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -155,23 +155,23 @@ fi
 
 
 
-(
-    . $MILABENCH_WORDIR/env/bin/activate
-    pip install git+https://github.com/breuleux/voir.git
+# (
+#     . $MILABENCH_WORDIR/env/bin/activate
+#     pip install git+https://github.com/breuleux/voir.git
 
-    (
-        . $BENCHMARK_VENV/bin/activate
-        pip install git+https://github.com/breuleux/voir.git
-    )
+#     (
+#         . $BENCHMARK_VENV/bin/activate
+#         pip install git+https://github.com/breuleux/voir.git
+#     )
 
-)
+# )
 
 # (
 #     # . $BENCHMARK_VENV/bin/activate
 #     # pip install xformers --index-url https://download.pytorch.org/whl/rocm6.1
 # )
 
-milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
+# milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
 
 # milabench prepare $ARGS # --system $MILABENCH_WORDIR/system.yaml
 # install_prepare

From 6c90b5efab6928b0df441bd5988b55d50b06adfc Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Thu, 27 Feb 2025 19:36:28 +0000
Subject: [PATCH 6/7] update voir

---
 .pin/constraints-cuda-torch.txt                     |  2 +-
 .pin/constraints-hpu-torch.txt                      |  2 +-
 .pin/constraints-rocm-torch.txt                     |  2 +-
 .pin/constraints-xpu-torch.txt                      |  2 +-
 benchmarks/accelerate_opt/requirements.cuda.txt     |  2 +-
 benchmarks/accelerate_opt/requirements.hpu.txt      |  2 +-
 benchmarks/accelerate_opt/requirements.rocm.txt     |  2 +-
 benchmarks/accelerate_opt/requirements.xpu.txt      |  2 +-
 benchmarks/brax/requirements.cuda.txt               |  2 +-
 benchmarks/brax/requirements.hpu.txt                |  2 +-
 benchmarks/brax/requirements.rocm.txt               |  2 +-
 benchmarks/brax/requirements.xpu.txt                |  2 +-
 benchmarks/diffusion/requirements.cuda.txt          |  2 +-
 benchmarks/diffusion/requirements.hpu.txt           |  2 +-
 benchmarks/diffusion/requirements.rocm.txt          |  2 +-
 benchmarks/diffusion/requirements.xpu.txt           |  2 +-
 benchmarks/dinov2/requirements.cuda.txt             |  2 +-
 benchmarks/dinov2/requirements.hpu.txt              |  2 +-
 benchmarks/dinov2/requirements.rocm.txt             |  2 +-
 benchmarks/dinov2/requirements.xpu.txt              |  2 +-
 benchmarks/dlrm/requirements.cuda.txt               |  2 +-
 benchmarks/dlrm/requirements.hpu.txt                |  2 +-
 benchmarks/dlrm/requirements.rocm.txt               |  2 +-
 benchmarks/dlrm/requirements.xpu.txt                |  2 +-
 benchmarks/flops/requirements.cuda.txt              |  2 +-
 benchmarks/flops/requirements.hpu.txt               |  2 +-
 benchmarks/flops/requirements.rocm.txt              |  2 +-
 benchmarks/flops/requirements.xpu.txt               |  2 +-
 benchmarks/geo_gnn/requirements.cuda.txt            |  2 +-
 benchmarks/geo_gnn/requirements.hpu.txt             |  2 +-
 benchmarks/geo_gnn/requirements.rocm.txt            |  2 +-
 benchmarks/huggingface/requirements.cuda.txt        |  2 +-
 benchmarks/huggingface/requirements.hpu.txt         |  2 +-
 benchmarks/huggingface/requirements.rocm.txt        |  2 +-
 benchmarks/huggingface/requirements.xpu.txt         |  2 +-
 benchmarks/lightning/requirements.cuda.txt          |  2 +-
 benchmarks/lightning/requirements.hpu.txt           |  2 +-
 benchmarks/lightning/requirements.rocm.txt          |  2 +-
 benchmarks/lightning/requirements.xpu.txt           |  2 +-
 benchmarks/llama/requirements.cuda.txt              |  2 +-
 benchmarks/llama/requirements.hpu.txt               |  2 +-
 benchmarks/llama/requirements.rocm.txt              |  2 +-
 benchmarks/llama/requirements.xpu.txt               |  2 +-
 benchmarks/llava/requirements.cuda.txt              |  2 +-
 benchmarks/llava/requirements.hpu.txt               |  2 +-
 benchmarks/llava/requirements.rocm.txt              |  2 +-
 benchmarks/llm/requirements.cuda.txt                |  2 +-
 benchmarks/llm/requirements.hpu.txt                 |  2 +-
 benchmarks/llm/requirements.rocm.txt                |  2 +-
 benchmarks/purejaxrl/requirements.cuda.txt          |  2 +-
 benchmarks/purejaxrl/requirements.hpu.txt           |  2 +-
 benchmarks/purejaxrl/requirements.rocm.txt          |  2 +-
 benchmarks/recursiongfn/requirements.cuda.txt       |  2 +-
 benchmarks/recursiongfn/requirements.hpu.txt        |  2 +-
 benchmarks/recursiongfn/requirements.rocm.txt       |  2 +-
 benchmarks/retired/rwkv/requirements.cuda.txt       |  2 +-
 benchmarks/retired/rwkv/requirements.hpu.txt        |  2 +-
 benchmarks/retired/rwkv/requirements.rocm.txt       |  2 +-
 benchmarks/retired/rwkv/requirements.xpu.txt        |  2 +-
 benchmarks/retired/stargan/requirements.cuda.txt    |  2 +-
 benchmarks/retired/stargan/requirements.hpu.txt     |  2 +-
 benchmarks/retired/stargan/requirements.rocm.txt    |  2 +-
 benchmarks/retired/stargan/requirements.xpu.txt     |  2 +-
 .../retired/super-slomo/requirements.cuda.txt       |  2 +-
 benchmarks/retired/super-slomo/requirements.hpu.txt |  2 +-
 .../retired/super-slomo/requirements.rocm.txt       |  2 +-
 benchmarks/retired/super-slomo/requirements.xpu.txt |  2 +-
 benchmarks/rlhf/requirements.cuda.txt               |  2 +-
 benchmarks/rlhf/requirements.hpu.txt                |  2 +-
 benchmarks/rlhf/requirements.rocm.txt               |  2 +-
 benchmarks/timm/requirements.cuda.txt               |  2 +-
 benchmarks/timm/requirements.hpu.txt                |  2 +-
 benchmarks/timm/requirements.rocm.txt               |  2 +-
 benchmarks/timm/requirements.xpu.txt                |  2 +-
 benchmarks/torchatari/requirements.cuda.txt         |  2 +-
 benchmarks/torchatari/requirements.hpu.txt          |  2 +-
 benchmarks/torchatari/requirements.rocm.txt         |  2 +-
 benchmarks/torchvision/requirements.cuda.txt        |  2 +-
 benchmarks/torchvision/requirements.hpu.txt         |  2 +-
 benchmarks/torchvision/requirements.rocm.txt        |  2 +-
 benchmarks/torchvision/requirements.xpu.txt         |  2 +-
 benchmarks/torchvision_ddp/requirements.cuda.txt    |  2 +-
 benchmarks/torchvision_ddp/requirements.hpu.txt     |  2 +-
 benchmarks/torchvision_ddp/requirements.rocm.txt    |  2 +-
 benchmarks/torchvision_ddp/requirements.xpu.txt     |  2 +-
 benchmarks/vjepa/requirements.cuda.txt              |  2 +-
 benchmarks/vjepa/requirements.hpu.txt               |  2 +-
 benchmarks/vjepa/requirements.rocm.txt              |  2 +-
 benchmate/benchmate/warden.py                       | 13 ++++++++++++-
 89 files changed, 100 insertions(+), 89 deletions(-)

diff --git a/.pin/constraints-cuda-torch.txt b/.pin/constraints-cuda-torch.txt
index 1c3b0208..16b92ade 100644
--- a/.pin/constraints-cuda-torch.txt
+++ b/.pin/constraints-cuda-torch.txt
@@ -990,7 +990,7 @@ urllib3==2.3.0
     #   sentry-sdk
 varname==0.14.0
     # via giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../constraints/cuda.txt
     #   -r benchmarks/brax/requirements.in
diff --git a/.pin/constraints-hpu-torch.txt b/.pin/constraints-hpu-torch.txt
index 92a55858..d1c40ee6 100644
--- a/.pin/constraints-hpu-torch.txt
+++ b/.pin/constraints-hpu-torch.txt
@@ -950,7 +950,7 @@ urllib3==2.2.3
     #   sentry-sdk
 varname==0.13.3
     # via giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../constraints/hpu.txt
     #   -r benchmarks/brax/requirements.in
diff --git a/.pin/constraints-rocm-torch.txt b/.pin/constraints-rocm-torch.txt
index cab2f08a..6e237dbc 100644
--- a/.pin/constraints-rocm-torch.txt
+++ b/.pin/constraints-rocm-torch.txt
@@ -946,7 +946,7 @@ urllib3==2.3.0
     #   sentry-sdk
 varname==0.14.0
     # via giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/constraints/rocm.txt
     #   -r benchmarks/brax/requirements.in
diff --git a/.pin/constraints-xpu-torch.txt b/.pin/constraints-xpu-torch.txt
index 9e427639..93064914 100644
--- a/.pin/constraints-xpu-torch.txt
+++ b/.pin/constraints-xpu-torch.txt
@@ -543,7 +543,7 @@ urllib3==2.2.2
     #   requests
 varname==0.10.0
     # via giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../constraints/xpu.txt
     #   -r benchmarks/brax/requirements.in
diff --git a/benchmarks/accelerate_opt/requirements.cuda.txt b/benchmarks/accelerate_opt/requirements.cuda.txt
index 52daee9b..f0b091e4 100644
--- a/benchmarks/accelerate_opt/requirements.cuda.txt
+++ b/benchmarks/accelerate_opt/requirements.cuda.txt
@@ -394,7 +394,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/accelerate_opt/requirements.hpu.txt b/benchmarks/accelerate_opt/requirements.hpu.txt
index 6af2b9b5..86c3b432 100644
--- a/benchmarks/accelerate_opt/requirements.hpu.txt
+++ b/benchmarks/accelerate_opt/requirements.hpu.txt
@@ -393,7 +393,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/accelerate_opt/requirements.rocm.txt b/benchmarks/accelerate_opt/requirements.rocm.txt
index 395bbd12..75a294c0 100644
--- a/benchmarks/accelerate_opt/requirements.rocm.txt
+++ b/benchmarks/accelerate_opt/requirements.rocm.txt
@@ -342,7 +342,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   -c .pin/../constraints/rocm.txt
diff --git a/benchmarks/accelerate_opt/requirements.xpu.txt b/benchmarks/accelerate_opt/requirements.xpu.txt
index e52b51c4..890496c9 100644
--- a/benchmarks/accelerate_opt/requirements.xpu.txt
+++ b/benchmarks/accelerate_opt/requirements.xpu.txt
@@ -341,7 +341,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/brax/requirements.cuda.txt b/benchmarks/brax/requirements.cuda.txt
index 8ced94d0..d9856f7d 100644
--- a/benchmarks/brax/requirements.cuda.txt
+++ b/benchmarks/brax/requirements.cuda.txt
@@ -439,7 +439,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/brax/requirements.hpu.txt b/benchmarks/brax/requirements.hpu.txt
index b02ff745..4dc7c258 100644
--- a/benchmarks/brax/requirements.hpu.txt
+++ b/benchmarks/brax/requirements.hpu.txt
@@ -410,7 +410,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/brax/requirements.rocm.txt b/benchmarks/brax/requirements.rocm.txt
index aed701cf..11f649c2 100644
--- a/benchmarks/brax/requirements.rocm.txt
+++ b/benchmarks/brax/requirements.rocm.txt
@@ -368,7 +368,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/brax/requirements.xpu.txt b/benchmarks/brax/requirements.xpu.txt
index 5e7dbe29..6614e420 100644
--- a/benchmarks/brax/requirements.xpu.txt
+++ b/benchmarks/brax/requirements.xpu.txt
@@ -355,7 +355,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/diffusion/requirements.cuda.txt b/benchmarks/diffusion/requirements.cuda.txt
index a84caa37..f4d594a7 100644
--- a/benchmarks/diffusion/requirements.cuda.txt
+++ b/benchmarks/diffusion/requirements.cuda.txt
@@ -420,7 +420,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/diffusion/requirements.hpu.txt b/benchmarks/diffusion/requirements.hpu.txt
index 88ccd569..2201eb08 100644
--- a/benchmarks/diffusion/requirements.hpu.txt
+++ b/benchmarks/diffusion/requirements.hpu.txt
@@ -362,7 +362,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/diffusion/requirements.rocm.txt b/benchmarks/diffusion/requirements.rocm.txt
index f73066f6..48489fe8 100644
--- a/benchmarks/diffusion/requirements.rocm.txt
+++ b/benchmarks/diffusion/requirements.rocm.txt
@@ -313,7 +313,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/diffusion/requirements.xpu.txt b/benchmarks/diffusion/requirements.xpu.txt
index 62a1aba1..aff0c0c7 100644
--- a/benchmarks/diffusion/requirements.xpu.txt
+++ b/benchmarks/diffusion/requirements.xpu.txt
@@ -313,7 +313,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.17
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/dinov2/requirements.cuda.txt b/benchmarks/dinov2/requirements.cuda.txt
index 42f002e6..d19eae85 100644
--- a/benchmarks/dinov2/requirements.cuda.txt
+++ b/benchmarks/dinov2/requirements.cuda.txt
@@ -295,7 +295,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/dinov2/requirements.hpu.txt b/benchmarks/dinov2/requirements.hpu.txt
index 4a11ccfb..b85cd164 100644
--- a/benchmarks/dinov2/requirements.hpu.txt
+++ b/benchmarks/dinov2/requirements.hpu.txt
@@ -249,7 +249,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/dinov2/requirements.rocm.txt b/benchmarks/dinov2/requirements.rocm.txt
index abf1fb06..a7914c2d 100644
--- a/benchmarks/dinov2/requirements.rocm.txt
+++ b/benchmarks/dinov2/requirements.rocm.txt
@@ -194,7 +194,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/dinov2/requirements.xpu.txt b/benchmarks/dinov2/requirements.xpu.txt
index 032296c6..f2a5ba0c 100644
--- a/benchmarks/dinov2/requirements.xpu.txt
+++ b/benchmarks/dinov2/requirements.xpu.txt
@@ -198,7 +198,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.17
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/dlrm/requirements.cuda.txt b/benchmarks/dlrm/requirements.cuda.txt
index b79010cf..f8b7f591 100644
--- a/benchmarks/dlrm/requirements.cuda.txt
+++ b/benchmarks/dlrm/requirements.cuda.txt
@@ -341,7 +341,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/dlrm/requirements.hpu.txt b/benchmarks/dlrm/requirements.hpu.txt
index d117a026..94d991b9 100644
--- a/benchmarks/dlrm/requirements.hpu.txt
+++ b/benchmarks/dlrm/requirements.hpu.txt
@@ -340,7 +340,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/dlrm/requirements.rocm.txt b/benchmarks/dlrm/requirements.rocm.txt
index 8debade4..7a9b92f3 100644
--- a/benchmarks/dlrm/requirements.rocm.txt
+++ b/benchmarks/dlrm/requirements.rocm.txt
@@ -289,7 +289,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   -c .pin/../constraints/rocm.txt
diff --git a/benchmarks/dlrm/requirements.xpu.txt b/benchmarks/dlrm/requirements.xpu.txt
index c263b4f2..f0d6ec85 100644
--- a/benchmarks/dlrm/requirements.xpu.txt
+++ b/benchmarks/dlrm/requirements.xpu.txt
@@ -285,7 +285,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/flops/requirements.cuda.txt b/benchmarks/flops/requirements.cuda.txt
index dc9bd142..2a67cb49 100644
--- a/benchmarks/flops/requirements.cuda.txt
+++ b/benchmarks/flops/requirements.cuda.txt
@@ -250,7 +250,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/flops/requirements.hpu.txt b/benchmarks/flops/requirements.hpu.txt
index 91e5677f..8feecdea 100644
--- a/benchmarks/flops/requirements.hpu.txt
+++ b/benchmarks/flops/requirements.hpu.txt
@@ -198,7 +198,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/flops/requirements.rocm.txt b/benchmarks/flops/requirements.rocm.txt
index 8297d303..3a116ee0 100644
--- a/benchmarks/flops/requirements.rocm.txt
+++ b/benchmarks/flops/requirements.rocm.txt
@@ -143,7 +143,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/flops/requirements.xpu.txt b/benchmarks/flops/requirements.xpu.txt
index 087e29b9..0dc17450 100644
--- a/benchmarks/flops/requirements.xpu.txt
+++ b/benchmarks/flops/requirements.xpu.txt
@@ -148,7 +148,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/geo_gnn/requirements.cuda.txt b/benchmarks/geo_gnn/requirements.cuda.txt
index 9cd54346..c209185a 100644
--- a/benchmarks/geo_gnn/requirements.cuda.txt
+++ b/benchmarks/geo_gnn/requirements.cuda.txt
@@ -365,7 +365,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/geo_gnn/requirements.hpu.txt b/benchmarks/geo_gnn/requirements.hpu.txt
index 9c6bb6d6..a6e6e570 100644
--- a/benchmarks/geo_gnn/requirements.hpu.txt
+++ b/benchmarks/geo_gnn/requirements.hpu.txt
@@ -310,7 +310,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/geo_gnn/requirements.rocm.txt b/benchmarks/geo_gnn/requirements.rocm.txt
index 41b911dd..47f3ca6f 100644
--- a/benchmarks/geo_gnn/requirements.rocm.txt
+++ b/benchmarks/geo_gnn/requirements.rocm.txt
@@ -251,7 +251,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/huggingface/requirements.cuda.txt b/benchmarks/huggingface/requirements.cuda.txt
index ad775694..72366e1c 100644
--- a/benchmarks/huggingface/requirements.cuda.txt
+++ b/benchmarks/huggingface/requirements.cuda.txt
@@ -290,7 +290,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/huggingface/requirements.hpu.txt b/benchmarks/huggingface/requirements.hpu.txt
index b5e21d99..cc06e0c3 100644
--- a/benchmarks/huggingface/requirements.hpu.txt
+++ b/benchmarks/huggingface/requirements.hpu.txt
@@ -239,7 +239,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/huggingface/requirements.rocm.txt b/benchmarks/huggingface/requirements.rocm.txt
index 684d3e4e..1d84fabe 100644
--- a/benchmarks/huggingface/requirements.rocm.txt
+++ b/benchmarks/huggingface/requirements.rocm.txt
@@ -183,7 +183,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/huggingface/requirements.xpu.txt b/benchmarks/huggingface/requirements.xpu.txt
index cb14810e..0f0552cf 100644
--- a/benchmarks/huggingface/requirements.xpu.txt
+++ b/benchmarks/huggingface/requirements.xpu.txt
@@ -185,7 +185,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/lightning/requirements.cuda.txt b/benchmarks/lightning/requirements.cuda.txt
index 08511440..5905245f 100644
--- a/benchmarks/lightning/requirements.cuda.txt
+++ b/benchmarks/lightning/requirements.cuda.txt
@@ -328,7 +328,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/lightning/requirements.hpu.txt b/benchmarks/lightning/requirements.hpu.txt
index f86fb064..b528dc2c 100644
--- a/benchmarks/lightning/requirements.hpu.txt
+++ b/benchmarks/lightning/requirements.hpu.txt
@@ -271,7 +271,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/lightning/requirements.rocm.txt b/benchmarks/lightning/requirements.rocm.txt
index 80b9558c..0367f94a 100644
--- a/benchmarks/lightning/requirements.rocm.txt
+++ b/benchmarks/lightning/requirements.rocm.txt
@@ -221,7 +221,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/lightning/requirements.xpu.txt b/benchmarks/lightning/requirements.xpu.txt
index 338ee0fb..ab92bf09 100644
--- a/benchmarks/lightning/requirements.xpu.txt
+++ b/benchmarks/lightning/requirements.xpu.txt
@@ -220,7 +220,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.17
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/llama/requirements.cuda.txt b/benchmarks/llama/requirements.cuda.txt
index 8cd9db85..949fa988 100644
--- a/benchmarks/llama/requirements.cuda.txt
+++ b/benchmarks/llama/requirements.cuda.txt
@@ -388,7 +388,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/llama/requirements.hpu.txt b/benchmarks/llama/requirements.hpu.txt
index 9c01a4dd..24e53fa4 100644
--- a/benchmarks/llama/requirements.hpu.txt
+++ b/benchmarks/llama/requirements.hpu.txt
@@ -330,7 +330,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/llama/requirements.rocm.txt b/benchmarks/llama/requirements.rocm.txt
index f43b4da1..814b9188 100644
--- a/benchmarks/llama/requirements.rocm.txt
+++ b/benchmarks/llama/requirements.rocm.txt
@@ -281,7 +281,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/llama/requirements.xpu.txt b/benchmarks/llama/requirements.xpu.txt
index b852606f..72378a2f 100644
--- a/benchmarks/llama/requirements.xpu.txt
+++ b/benchmarks/llama/requirements.xpu.txt
@@ -280,7 +280,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/llava/requirements.cuda.txt b/benchmarks/llava/requirements.cuda.txt
index a437d671..231520fa 100644
--- a/benchmarks/llava/requirements.cuda.txt
+++ b/benchmarks/llava/requirements.cuda.txt
@@ -386,7 +386,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/llava/requirements.hpu.txt b/benchmarks/llava/requirements.hpu.txt
index 3bd40dff..2db7dbcb 100644
--- a/benchmarks/llava/requirements.hpu.txt
+++ b/benchmarks/llava/requirements.hpu.txt
@@ -328,7 +328,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/llava/requirements.rocm.txt b/benchmarks/llava/requirements.rocm.txt
index e6d94665..ca293d0a 100644
--- a/benchmarks/llava/requirements.rocm.txt
+++ b/benchmarks/llava/requirements.rocm.txt
@@ -279,7 +279,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/llm/requirements.cuda.txt b/benchmarks/llm/requirements.cuda.txt
index 5bf56ae8..215cbfe8 100644
--- a/benchmarks/llm/requirements.cuda.txt
+++ b/benchmarks/llm/requirements.cuda.txt
@@ -460,7 +460,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/llm/requirements.hpu.txt b/benchmarks/llm/requirements.hpu.txt
index 9b88be53..55c5834a 100644
--- a/benchmarks/llm/requirements.hpu.txt
+++ b/benchmarks/llm/requirements.hpu.txt
@@ -393,7 +393,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/llm/requirements.rocm.txt b/benchmarks/llm/requirements.rocm.txt
index 29c3aab3..c645ea8d 100644
--- a/benchmarks/llm/requirements.rocm.txt
+++ b/benchmarks/llm/requirements.rocm.txt
@@ -358,7 +358,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/purejaxrl/requirements.cuda.txt b/benchmarks/purejaxrl/requirements.cuda.txt
index 416807f4..f53f80ef 100644
--- a/benchmarks/purejaxrl/requirements.cuda.txt
+++ b/benchmarks/purejaxrl/requirements.cuda.txt
@@ -770,7 +770,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/purejaxrl/requirements.hpu.txt b/benchmarks/purejaxrl/requirements.hpu.txt
index aeb2b110..5926b4ba 100644
--- a/benchmarks/purejaxrl/requirements.hpu.txt
+++ b/benchmarks/purejaxrl/requirements.hpu.txt
@@ -721,7 +721,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/purejaxrl/requirements.rocm.txt b/benchmarks/purejaxrl/requirements.rocm.txt
index dcd78c0e..626b56fb 100644
--- a/benchmarks/purejaxrl/requirements.rocm.txt
+++ b/benchmarks/purejaxrl/requirements.rocm.txt
@@ -699,7 +699,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/recursiongfn/requirements.cuda.txt b/benchmarks/recursiongfn/requirements.cuda.txt
index af2df13b..0373a101 100644
--- a/benchmarks/recursiongfn/requirements.cuda.txt
+++ b/benchmarks/recursiongfn/requirements.cuda.txt
@@ -538,7 +538,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/recursiongfn/requirements.hpu.txt b/benchmarks/recursiongfn/requirements.hpu.txt
index 4e362ae6..7c399f9c 100644
--- a/benchmarks/recursiongfn/requirements.hpu.txt
+++ b/benchmarks/recursiongfn/requirements.hpu.txt
@@ -471,7 +471,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/recursiongfn/requirements.rocm.txt b/benchmarks/recursiongfn/requirements.rocm.txt
index ec19d7cb..c0ae1d79 100644
--- a/benchmarks/recursiongfn/requirements.rocm.txt
+++ b/benchmarks/recursiongfn/requirements.rocm.txt
@@ -445,7 +445,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/retired/rwkv/requirements.cuda.txt b/benchmarks/retired/rwkv/requirements.cuda.txt
index d9c29cbb..3e38a391 100644
--- a/benchmarks/retired/rwkv/requirements.cuda.txt
+++ b/benchmarks/retired/rwkv/requirements.cuda.txt
@@ -266,7 +266,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via -r benchmarks/rwkv/requirements.in
 yarl==1.9.2
     # via
diff --git a/benchmarks/retired/rwkv/requirements.hpu.txt b/benchmarks/retired/rwkv/requirements.hpu.txt
index 4fbe3260..e147c1e9 100644
--- a/benchmarks/retired/rwkv/requirements.hpu.txt
+++ b/benchmarks/retired/rwkv/requirements.hpu.txt
@@ -265,7 +265,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via -r milabench/benchmarks/rwkv/requirements.in
 yarl==1.9.4
     # via
diff --git a/benchmarks/retired/rwkv/requirements.rocm.txt b/benchmarks/retired/rwkv/requirements.rocm.txt
index e3ca029d..60c67e07 100644
--- a/benchmarks/retired/rwkv/requirements.rocm.txt
+++ b/benchmarks/retired/rwkv/requirements.rocm.txt
@@ -214,7 +214,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via -r benchmarks/rwkv/requirements.in
 yarl==1.9.2
     # via
diff --git a/benchmarks/retired/rwkv/requirements.xpu.txt b/benchmarks/retired/rwkv/requirements.xpu.txt
index 4dab5ace..94275598 100644
--- a/benchmarks/retired/rwkv/requirements.xpu.txt
+++ b/benchmarks/retired/rwkv/requirements.xpu.txt
@@ -267,7 +267,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/retired/stargan/requirements.cuda.txt b/benchmarks/retired/stargan/requirements.cuda.txt
index 749cfacd..a67296fe 100644
--- a/benchmarks/retired/stargan/requirements.cuda.txt
+++ b/benchmarks/retired/stargan/requirements.cuda.txt
@@ -231,7 +231,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/retired/stargan/requirements.hpu.txt b/benchmarks/retired/stargan/requirements.hpu.txt
index 11862711..9b758cdf 100644
--- a/benchmarks/retired/stargan/requirements.hpu.txt
+++ b/benchmarks/retired/stargan/requirements.hpu.txt
@@ -230,7 +230,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/retired/stargan/requirements.rocm.txt b/benchmarks/retired/stargan/requirements.rocm.txt
index c5fa7a2c..48a808ff 100644
--- a/benchmarks/retired/stargan/requirements.rocm.txt
+++ b/benchmarks/retired/stargan/requirements.rocm.txt
@@ -179,7 +179,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   -c .pin/../constraints/rocm.txt
diff --git a/benchmarks/retired/stargan/requirements.xpu.txt b/benchmarks/retired/stargan/requirements.xpu.txt
index 7246fd7e..35efaf2a 100644
--- a/benchmarks/retired/stargan/requirements.xpu.txt
+++ b/benchmarks/retired/stargan/requirements.xpu.txt
@@ -177,7 +177,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/retired/super-slomo/requirements.cuda.txt b/benchmarks/retired/super-slomo/requirements.cuda.txt
index 66ce0258..6ca81c8f 100644
--- a/benchmarks/retired/super-slomo/requirements.cuda.txt
+++ b/benchmarks/retired/super-slomo/requirements.cuda.txt
@@ -249,7 +249,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/retired/super-slomo/requirements.hpu.txt b/benchmarks/retired/super-slomo/requirements.hpu.txt
index a386b0e4..eda4be36 100644
--- a/benchmarks/retired/super-slomo/requirements.hpu.txt
+++ b/benchmarks/retired/super-slomo/requirements.hpu.txt
@@ -198,7 +198,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/retired/super-slomo/requirements.rocm.txt b/benchmarks/retired/super-slomo/requirements.rocm.txt
index d85fcf14..84b400ac 100644
--- a/benchmarks/retired/super-slomo/requirements.rocm.txt
+++ b/benchmarks/retired/super-slomo/requirements.rocm.txt
@@ -144,7 +144,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-rocm-torch.txt
     #   -c .pin/../constraints/rocm.txt
diff --git a/benchmarks/retired/super-slomo/requirements.xpu.txt b/benchmarks/retired/super-slomo/requirements.xpu.txt
index 1a40b14f..60b1afcf 100644
--- a/benchmarks/retired/super-slomo/requirements.xpu.txt
+++ b/benchmarks/retired/super-slomo/requirements.xpu.txt
@@ -144,7 +144,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/extra/torch.xpu.txt
diff --git a/benchmarks/rlhf/requirements.cuda.txt b/benchmarks/rlhf/requirements.cuda.txt
index b01faab9..c13d6ec5 100644
--- a/benchmarks/rlhf/requirements.cuda.txt
+++ b/benchmarks/rlhf/requirements.cuda.txt
@@ -393,7 +393,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/rlhf/requirements.hpu.txt b/benchmarks/rlhf/requirements.hpu.txt
index 1c1bcce6..ccd8184d 100644
--- a/benchmarks/rlhf/requirements.hpu.txt
+++ b/benchmarks/rlhf/requirements.hpu.txt
@@ -347,7 +347,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/rlhf/requirements.rocm.txt b/benchmarks/rlhf/requirements.rocm.txt
index 9e91a184..fc6d873c 100644
--- a/benchmarks/rlhf/requirements.rocm.txt
+++ b/benchmarks/rlhf/requirements.rocm.txt
@@ -286,7 +286,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/timm/requirements.cuda.txt b/benchmarks/timm/requirements.cuda.txt
index 14635500..3f25fb4c 100644
--- a/benchmarks/timm/requirements.cuda.txt
+++ b/benchmarks/timm/requirements.cuda.txt
@@ -278,7 +278,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/timm/requirements.hpu.txt b/benchmarks/timm/requirements.hpu.txt
index e626bd1f..f919e8b7 100644
--- a/benchmarks/timm/requirements.hpu.txt
+++ b/benchmarks/timm/requirements.hpu.txt
@@ -226,7 +226,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/timm/requirements.rocm.txt b/benchmarks/timm/requirements.rocm.txt
index c493a3f6..dd1463da 100644
--- a/benchmarks/timm/requirements.rocm.txt
+++ b/benchmarks/timm/requirements.rocm.txt
@@ -171,7 +171,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/timm/requirements.xpu.txt b/benchmarks/timm/requirements.xpu.txt
index d7f99337..07d57348 100644
--- a/benchmarks/timm/requirements.xpu.txt
+++ b/benchmarks/timm/requirements.xpu.txt
@@ -174,7 +174,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/torchatari/requirements.cuda.txt b/benchmarks/torchatari/requirements.cuda.txt
index 90f19eb1..c19f8638 100644
--- a/benchmarks/torchatari/requirements.cuda.txt
+++ b/benchmarks/torchatari/requirements.cuda.txt
@@ -350,7 +350,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/torchatari/requirements.hpu.txt b/benchmarks/torchatari/requirements.hpu.txt
index 6d7369df..90760a4e 100644
--- a/benchmarks/torchatari/requirements.hpu.txt
+++ b/benchmarks/torchatari/requirements.hpu.txt
@@ -290,7 +290,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/torchatari/requirements.rocm.txt b/benchmarks/torchatari/requirements.rocm.txt
index c01583b3..21a86a0b 100644
--- a/benchmarks/torchatari/requirements.rocm.txt
+++ b/benchmarks/torchatari/requirements.rocm.txt
@@ -249,7 +249,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/torchvision/requirements.cuda.txt b/benchmarks/torchvision/requirements.cuda.txt
index 4090e7b9..f92d305b 100644
--- a/benchmarks/torchvision/requirements.cuda.txt
+++ b/benchmarks/torchvision/requirements.cuda.txt
@@ -250,7 +250,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/torchvision/requirements.hpu.txt b/benchmarks/torchvision/requirements.hpu.txt
index f0b47e91..66840493 100644
--- a/benchmarks/torchvision/requirements.hpu.txt
+++ b/benchmarks/torchvision/requirements.hpu.txt
@@ -198,7 +198,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/torchvision/requirements.rocm.txt b/benchmarks/torchvision/requirements.rocm.txt
index f9d737fb..184700b9 100644
--- a/benchmarks/torchvision/requirements.rocm.txt
+++ b/benchmarks/torchvision/requirements.rocm.txt
@@ -143,7 +143,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/torchvision/requirements.xpu.txt b/benchmarks/torchvision/requirements.xpu.txt
index 3cd87697..472dc26f 100644
--- a/benchmarks/torchvision/requirements.xpu.txt
+++ b/benchmarks/torchvision/requirements.xpu.txt
@@ -146,7 +146,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/torchvision_ddp/requirements.cuda.txt b/benchmarks/torchvision_ddp/requirements.cuda.txt
index 60a3f9f8..305c9f2e 100644
--- a/benchmarks/torchvision_ddp/requirements.cuda.txt
+++ b/benchmarks/torchvision_ddp/requirements.cuda.txt
@@ -250,7 +250,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/torchvision_ddp/requirements.hpu.txt b/benchmarks/torchvision_ddp/requirements.hpu.txt
index a4174e7b..3cadb0aa 100644
--- a/benchmarks/torchvision_ddp/requirements.hpu.txt
+++ b/benchmarks/torchvision_ddp/requirements.hpu.txt
@@ -198,7 +198,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/torchvision_ddp/requirements.rocm.txt b/benchmarks/torchvision_ddp/requirements.rocm.txt
index 22400c72..03102cc1 100644
--- a/benchmarks/torchvision_ddp/requirements.rocm.txt
+++ b/benchmarks/torchvision_ddp/requirements.rocm.txt
@@ -143,7 +143,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmarks/torchvision_ddp/requirements.xpu.txt b/benchmarks/torchvision_ddp/requirements.xpu.txt
index b3f732e8..29bfe96c 100644
--- a/benchmarks/torchvision_ddp/requirements.xpu.txt
+++ b/benchmarks/torchvision_ddp/requirements.xpu.txt
@@ -145,7 +145,7 @@ varname==0.10.0
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-xpu-torch.txt
     #   -c .pin/../constraints/xpu.txt
diff --git a/benchmarks/vjepa/requirements.cuda.txt b/benchmarks/vjepa/requirements.cuda.txt
index 9dcbe4b3..2ac61fdb 100644
--- a/benchmarks/vjepa/requirements.cuda.txt
+++ b/benchmarks/vjepa/requirements.cuda.txt
@@ -341,7 +341,7 @@ varname==0.14.0
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-cuda-torch.txt
     #   -c .pin/../constraints/cuda.txt
diff --git a/benchmarks/vjepa/requirements.hpu.txt b/benchmarks/vjepa/requirements.hpu.txt
index b1c986ec..ba184de3 100644
--- a/benchmarks/vjepa/requirements.hpu.txt
+++ b/benchmarks/vjepa/requirements.hpu.txt
@@ -286,7 +286,7 @@ varname==0.13.3
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c .pin/../.pin/constraints-hpu-torch.txt
     #   -c .pin/../constraints/hpu.txt
diff --git a/benchmarks/vjepa/requirements.rocm.txt b/benchmarks/vjepa/requirements.rocm.txt
index 0ffac47b..2a0d5405 100644
--- a/benchmarks/vjepa/requirements.rocm.txt
+++ b/benchmarks/vjepa/requirements.rocm.txt
@@ -234,7 +234,7 @@ varname==0.14.0
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   giving
-voir==0.2.19
+voir==0.2.20
     # via
     #   -c /home/testroot/milabench/.pin/constraints-rocm-torch.txt
     #   -c /home/testroot/milabench/constraints/rocm.txt
diff --git a/benchmate/benchmate/warden.py b/benchmate/benchmate/warden.py
index 7a1a9dbe..5b97f6a9 100644
--- a/benchmate/benchmate/warden.py
+++ b/benchmate/benchmate/warden.py
@@ -73,7 +73,7 @@ def _default():
 
 
 def _rocm_parse_processes():
-    cmd = ["rocm-smi", f"--showpids", "--json"]
+    cmd = ["rocm-smi", f"--showpids", "--json", "--loglevel", "error"]
     output = subprocess.check_output(cmd, text=True)
     
     info = []
@@ -213,6 +213,17 @@ def wait_for_children():
     for child in children:
         os.kill(child, signal.SIGCONT)
 
+    count = 0
+    for child in children:
+        try:
+            pid, status = os.waitpid(child, os.WNOHANG)
+            count += 1
+        except:
+            pass
+
+    if count > 0:
+        syslog(f"{count}/{len(children)} reaped children")
+
 
 class GPUProcessWarden(BaseWarden):
     """Ensure all the processes using the GPU are killed before & after the bench"""

From a848a86e2d731bae2b8b0fd7655a044455a35384 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Fri, 28 Feb 2025 02:52:11 +0000
Subject: [PATCH 7/7] Tweaks

---
 benchmate/benchmate/warden.py | 40 +++++++++++++++++------------------
 pyproject.toml                |  2 +-
 scripts/article/run_rocm.sh   | 19 +++++++++--------
 3 files changed, 31 insertions(+), 30 deletions(-)

diff --git a/benchmate/benchmate/warden.py b/benchmate/benchmate/warden.py
index 5b97f6a9..84b18880 100644
--- a/benchmate/benchmate/warden.py
+++ b/benchmate/benchmate/warden.py
@@ -71,9 +71,8 @@ def _default():
     return []
 
 
-
 def _rocm_parse_processes():
-    cmd = ["rocm-smi", f"--showpids", "--json", "--loglevel", "error"]
+    cmd = ["rocm-smi", "--loglevel", "error", "--showpids", "--json"]
     output = subprocess.check_output(cmd, text=True)
     
     info = []
@@ -84,7 +83,14 @@ def _rocm_parse_processes():
         return info
 
     for key, data in data.get("system", {}).items():
-        process_name, ngpu, vram, sdma, cu_occupancy = data.split(",")
+        cols = data.split(",")
+
+        process_name = cols[0]
+        ngpu = cols[1]
+        vram = cols[2]
+        sdma = cols[3]
+        cu_occupancy = cols[4]
+
         pid = key[3:]
         info.append(ProcessInfo(pid=pid, process_name=process_name, used_memory=vram))
     return info
@@ -167,6 +173,7 @@ def children_warden(enabled=True):
         yield
         return
 
+    pid = os.getpid()
 
     def get_children():
         with open(f"/proc/{pid}/task/{pid}/children", "r") as f:
@@ -174,25 +181,27 @@ def get_children():
 
     prev = set(get_children())
 
+    yield
+
     def get_children():
         with open(f"/proc/{pid}/task/{pid}/children", "r") as f:
             return set([int(c) for c in f.read().strip().split()]) - prev
 
-    yield
-
-    pid = os.getpid()
-
-    def wait_for_children():
+    def wait_for_children(wait_time=15):
         children = get_children()
         start_time = time.time()
 
-        while children and time.time()  - start_time < 15:
+        while children and time.time()  - start_time < wait_time:
+            for child in children:
+                os.waitpid(child, os.WNOHANG)
+
             children = get_children()
             time.sleep(0)
 
         return children
 
-    children = get_children()
+    children = wait_for_children(1)
+
     if children:
         syslog(f"{len(children)} still alive after end of benchmark")
 
@@ -213,16 +222,7 @@ def wait_for_children():
     for child in children:
         os.kill(child, signal.SIGCONT)
 
-    count = 0
-    for child in children:
-        try:
-            pid, status = os.waitpid(child, os.WNOHANG)
-            count += 1
-        except:
-            pass
-
-    if count > 0:
-        syslog(f"{count}/{len(children)} reaped children")
+    wait_for_children(1)
 
 
 class GPUProcessWarden(BaseWarden):
diff --git a/pyproject.toml b/pyproject.toml
index 7efb653d..272be9bf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,7 +10,7 @@ authors = [
 license = "MIT"
 
 [tool.poetry.dependencies]
-voir =  ">=0.2.19"
+voir =  ">=0.2.20"
 benchmate =  {path = "benchmate", develop = false}
 python = ">=3.10,<4.0"
 giving = "^0.4.0"
diff --git a/scripts/article/run_rocm.sh b/scripts/article/run_rocm.sh
index 85f1c8b3..4cd9ea87 100644
--- a/scripts/article/run_rocm.sh
+++ b/scripts/article/run_rocm.sh
@@ -155,16 +155,17 @@ fi
 
 
 
-# (
-#     . $MILABENCH_WORDIR/env/bin/activate
-#     pip install git+https://github.com/breuleux/voir.git
-
-#     (
-#         . $BENCHMARK_VENV/bin/activate
-#         pip install git+https://github.com/breuleux/voir.git
-#     )
+(
+    . $MILABENCH_WORDIR/env/bin/activate
+    pip install git+https://github.com/breuleux/voir.git
+    pip install -e /home/testroot/milabench/benchmate
+    (
+        . $BENCHMARK_VENV/bin/activate
+        pip install git+https://github.com/breuleux/voir.git
+        pip install -e /home/testroot/milabench/benchmate
+    )
 
-# )
+)
 
 # (
 #     # . $BENCHMARK_VENV/bin/activate