Fix merge error (#2595)

huggingface · Jan 20, 2025 · 3d2c1e4 · 3d2c1e4
1 parent 5fd7836
commit 3d2c1e4
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -44,6 +44,8 @@ def test_env(self, mock_stdout):
         self.assertIn("TRL version: ", mock_stdout.getvalue().strip())
 
     def test_grpo(self):
+        from trl.cli import main
+
         with tempfile.TemporaryDirectory() as tmp_dir:  # Create a temporary directory
             command = f"trl grpo --output_dir {tmp_dir} --model_name_or_path trl-internal-testing/tiny-Qwen2ForCausalLM-2.5 --reward_model_name_or_path trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5 --dataset_name trl-internal-testing/zen --dataset_config standard_prompt_only --num_generations 3 --max_completion_length 32 --report_to none"
             with patch("sys.argv", command.split(" ")):