diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 6d9ac0cc63..e10661f7c8 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -213,10 +213,9 @@ jobs:
       - if: ${{ matrix.task == 'onnx' }}
         name: Test ONNX (CLANG)
         run: CLANG=1 python -m pytest -n=auto test/external/external_test_onnx_backend.py --durations=20
-      # NOTE: need to regenerate the dataset for this after the reduce change. it should probably gen a small dataset here instead
-      #- if: ${{ matrix.task == 'onnx' }}
-      #  name: Test Action Space
-      #  run: PYTHONPATH="." GPU=1 python3 extra/optimization/get_action_space.py
+      - if: ${{ matrix.task == 'onnx' }}
+        name: Test Action Space
+        run: PYTHONPATH="." GPU=1 python3 extra/optimization/get_action_space.py
       - if: ${{ matrix.task == 'onnx' }}
         name: Test Beam Search
         run: PYTHONPATH="." GPU=1 IGNORE_BEAM_CACHE=1 python3 -m pytest extra/optimization/test_beam_search.py
diff --git a/extra/datasets/sops.gz b/extra/datasets/sops.gz
index 2edf22547c..3aa9d8fe5d 100644
Binary files a/extra/datasets/sops.gz and b/extra/datasets/sops.gz differ