From 7fb3aad53d087444e318b2a694208a5b3529ebda Mon Sep 17 00:00:00 2001 From: Artur Kloniecki Date: Fri, 17 Oct 2025 07:58:24 +0300 Subject: [PATCH] Enable Int64 support for trl-sft workload. Signed-off-by: Artur Kloniecki --- tests/test_examples.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test_examples.py b/tests/test_examples.py index 03a20cd251..99f79783d4 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -532,6 +532,9 @@ def test(self): if "--use_hpu_graphs_for_inference" in extra_command_line_arguments: extra_command_line_arguments.remove("--use_hpu_graphs_for_inference") + if self.TASK_NAME == "trl-sft": + env_variables["PT_ENABLE_INT64_SUPPORT"] = "1" + extra_command_line_arguments += self._get_dataset_args() if torch_compile and (