Minor fix for example tests

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

.github/workflows/example_tests.yml

@@ -74,7 +74,7 @@ jobs:
       example: ${{ matrix.example }}
       timeout_minutes: 30
       pip_install_extras: "[hf,dev-test]"
-      runner: linux-amd64-gpu-h100-latest-1
+      runner: linux-amd64-gpu-rtxpro6000-latest-1
 
   torch-non-pr:
     if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}
@@ -161,7 +161,7 @@ jobs:
       docker_image: "nvcr.io/nvidia/tensorrt:26.02-py3"
       example: ${{ matrix.example }}
       pip_install_extras: "[all,dev-test]"
-      runner: linux-amd64-gpu-l4-latest-1
+      runner: linux-amd64-gpu-rtxpro6000-latest-1
 
   onnx-non-pr:
     if: ${{ !startsWith(github.ref, 'refs/heads/pull-request/') }}

examples/llm_sparsity/attention_sparsity/hf_sa.py

@@ -111,7 +111,7 @@ def generate_sample_output(model, tokenizer, args):
         padding=False,
     )
     if torch.cuda.is_available():
-        inputs = {k: v.cuda() for k, v in inputs.items()}
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         # Generate
         with torch.no_grad():

examples/speculative_decoding/.gitignore

@@ -1,4 +1,5 @@
 Daring-Anteater
 input_conversations
 synthetic_conversations
-ckpts
+ckpts
+mlflow.db

tests/examples/speculative_decoding/conftest.py

@@ -36,8 +36,7 @@ def tiny_daring_anteater_path(tmp_path_factory):
     config_path.write_text(yaml.dump(config))
 
     run_example_command(
-        ["python", "../dataset/make_dataset.py", "-f", str(config_path), "--full-conversations"],
-        "speculative_decoding",
+        ["python", "make_dataset.py", "-f", str(config_path), "--full-conversations"], "dataset"
     )
 
     return output_file

tools/launcher/common/eagle3/make_dataset.sh

@@ -23,7 +23,7 @@ source ${SCRIPT_DIR}/../service_utils.sh
 
 trap 'error_handler $0 $LINENO' ERR # ERROR HANDLER
 
-python modules/Model-Optimizer/examples/speculative_decoding/prepare_input_conversations/make_dataset.py \
+python modules/Model-Optimizer/examples/dataset/make_dataset.py \
     ${@}
 
 mkdir -p /scratchspace/data