diff --git a/.github/workflows/build-wheels-aarch64-linux.yml b/.github/workflows/build-wheels-aarch64-linux.yml
index 33a3d0030f0..30509d1a6df 100644
--- a/.github/workflows/build-wheels-aarch64-linux.yml
+++ b/.github/workflows/build-wheels-aarch64-linux.yml
@@ -19,7 +19,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: vmoens/test-infra/.github/workflows/generate_binary_build_matrix.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@pull/6151
     with:
       package-type: wheel
       os: linux-aarch64
@@ -36,7 +36,7 @@ jobs:
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
     name: pytorch/rl
-    uses: vmoens/test-infra/.github/workflows/build_wheels_linux.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@pull/6151
     with:
       repository: ${{ matrix.repository }}
       ref: ""
diff --git a/.github/workflows/build-wheels-linux.yml b/.github/workflows/build-wheels-linux.yml
index fa1bf0323d7..18d7f850eb7 100644
--- a/.github/workflows/build-wheels-linux.yml
+++ b/.github/workflows/build-wheels-linux.yml
@@ -19,7 +19,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: vmoens/test-infra/.github/workflows/generate_binary_build_matrix.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@pull/6151
     with:
       package-type: wheel
       os: linux
@@ -35,7 +35,7 @@ jobs:
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
     name: pytorch/rl
-    uses: vmoens/test-infra/.github/workflows/build_wheels_linux.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@pull/6151
     with:
       repository: ${{ matrix.repository }}
       ref: ""
diff --git a/.github/workflows/build-wheels-m1.yml b/.github/workflows/build-wheels-m1.yml
index 12f136d9e7a..19de21a17cb 100644
--- a/.github/workflows/build-wheels-m1.yml
+++ b/.github/workflows/build-wheels-m1.yml
@@ -19,7 +19,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: vmoens/test-infra/.github/workflows/generate_binary_build_matrix.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@pull/6151
     with:
       package-type: wheel
       os: macos-arm64
@@ -35,7 +35,7 @@ jobs:
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
     name: pytorch/rl
-    uses: vmoens/test-infra/.github/workflows/build_wheels_macos.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/build_wheels_macos.yml@pull/6151
     with:
       repository: ${{ matrix.repository }}
       ref: ""
diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml
index 62a1138e66e..bc46c61738b 100644
--- a/.github/workflows/build-wheels-windows.yml
+++ b/.github/workflows/build-wheels-windows.yml
@@ -19,7 +19,7 @@ permissions:
 
 jobs:
   generate-matrix:
-    uses: vmoens/test-infra/.github/workflows/generate_binary_build_matrix.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@pull/6151
     with:
       package-type: wheel
       os: windows
@@ -38,7 +38,7 @@ jobs:
             smoke-test-script: test/smoke_test.py
             package-name: torchrl
     name: ${{ matrix.repository }}
-    uses: vmoens/test-infra/.github/workflows/build_wheels_windows.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/build_wheels_windows.yml@pull/6151
     with:
       repository: ${{ matrix.repository }}
       ref: ""
diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml
index 78d8265b4b9..2376b3cb1bd 100644
--- a/.github/workflows/docs.yml
+++ b/.github/workflows/docs.yml
@@ -24,7 +24,7 @@ jobs:
       matrix:
         python_version: ["3.10"]
         cuda_arch_version: ["12.1"]
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       upload-artifact: docs
@@ -118,7 +118,7 @@ jobs:
         ((github.ref_type == 'branch' && github.ref_name == 'main') || github.ref_type == 'tag')
     permissions:
       contents: write
-    uses: vmoens/test-infra/.github/workflows/linux_job.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job.yml@pull/6151
     with:
       repository: pytorch/rl
       download-artifact: docs
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 2f997e64c09..91e9a65b3b1 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   python-source-and-configs:
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       script: |
@@ -45,7 +45,7 @@ jobs:
         echo '::endgroup::'
 
   c-source:
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       script: |
diff --git a/.github/workflows/test-linux-habitat.yml b/.github/workflows/test-linux-habitat.yml
index 6d22017d783..e3933c4d98f 100644
--- a/.github/workflows/test-linux-habitat.yml
+++ b/.github/workflows/test-linux-habitat.yml
@@ -22,7 +22,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
diff --git a/.github/workflows/test-linux-libs.yml b/.github/workflows/test-linux-libs.yml
index 42afee05c45..93017174576 100644
--- a/.github/workflows/test-linux-libs.yml
+++ b/.github/workflows/test-linux-libs.yml
@@ -23,7 +23,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -57,7 +57,7 @@ jobs:
         python_version: ["3.11"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -94,7 +94,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -129,7 +129,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -164,7 +164,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -197,7 +197,7 @@ jobs:
       matrix:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -233,7 +233,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -267,7 +267,7 @@ jobs:
         bash .github/unittest/linux_libs/scripts_jumanji/post_process.sh
 
   unittests-meltingpot:
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
     with:
       repository: pytorch/rl
@@ -307,7 +307,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -345,7 +345,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -383,7 +383,7 @@ jobs:
         python_version: ["3.10.12"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -421,7 +421,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -456,7 +456,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -487,7 +487,7 @@ jobs:
 
   unittests-pettingzoo:
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -526,7 +526,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -560,7 +560,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -594,7 +594,7 @@ jobs:
       matrix:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -631,7 +631,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -670,7 +670,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Data') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -705,7 +705,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
     if: ${{ github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Environments') }}
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
diff --git a/.github/workflows/test-linux-rlhf.yml b/.github/workflows/test-linux-rlhf.yml
index b3a4f82a9d5..41394105e00 100644
--- a/.github/workflows/test-linux-rlhf.yml
+++ b/.github/workflows/test-linux-rlhf.yml
@@ -21,7 +21,7 @@ jobs:
       matrix:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
diff --git a/.github/workflows/test-linux-sota.yml b/.github/workflows/test-linux-sota.yml
index efd4d44542a..f3a520a9e09 100644
--- a/.github/workflows/test-linux-sota.yml
+++ b/.github/workflows/test-linux-sota.yml
@@ -25,7 +25,7 @@ jobs:
         python_version: ["3.9"]
         cuda_arch_version: ["12.1"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
diff --git a/.github/workflows/test-linux.yml b/.github/workflows/test-linux.yml
index b180421d2d3..47b8e7bfa8a 100644
--- a/.github/workflows/test-linux.yml
+++ b/.github/workflows/test-linux.yml
@@ -24,7 +24,7 @@ jobs:
       matrix:
         python_version: ["3.9", "3.10", "3.11", "3.12"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.12xlarge
       repository: pytorch/rl
@@ -55,7 +55,7 @@ jobs:
       matrix:
         python_version: ["3.12"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.12xlarge
       repository: pytorch/rl
@@ -87,7 +87,7 @@ jobs:
         python_version: ["3.11"]
         cuda_arch_version: ["12.1"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
@@ -124,7 +124,7 @@ jobs:
       matrix:
         python_version: ["3.8"]
         cuda_arch_version: ["11.6"]
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       repository: pytorch/rl
       runner: "linux.g5.4xlarge.nvidia.gpu"
@@ -156,7 +156,7 @@ jobs:
         python_version: ["3.11"]
         cuda_arch_version: ["12.1"]
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
@@ -191,7 +191,7 @@ jobs:
         python_version: ["3.10"] # "3.9", "3.10", "3.11"
         cuda_arch_version: ["11.8"] # "11.6", "11.7"
       fail-fast: false
-    uses: vmoens/test-infra/.github/workflows/linux_job_v2.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@pull/6151
     with:
       runner: linux.g5.4xlarge.nvidia.gpu
       repository: pytorch/rl
diff --git a/.github/workflows/test-windows-optdepts.yml b/.github/workflows/test-windows-optdepts.yml
index 210bfc6d2c8..96eb6de5b28 100644
--- a/.github/workflows/test-windows-optdepts.yml
+++ b/.github/workflows/test-windows-optdepts.yml
@@ -17,7 +17,7 @@ concurrency:
 
 jobs:
   unittests-cpu:
-    uses: vmoens/test-infra/.github/workflows/windows_job.yml@download-v4
+    uses: pytorch/test-infra/.github/workflows/windows_job.yml@pull/6151
     with:
       runner: "windows.4xlarge"
       repository: pytorch/rl
@@ -57,7 +57,7 @@ jobs:
         ./.github/unittest/windows_optdepts/scripts/post_process.sh
 
 #  unittests-gpu:
-#    uses: vmoens/test-infra/.github/workflows/windows_job.yml@download-v4
+#    uses: pytorch/test-infra/.github/workflows/windows_job.yml@pull/6151
 #    with:
 #      runner: "windows.g5.4xlarge.nvidia.gpu"
 #      repository: pytorch/rl