Fix failling GGML test (#34871)

fix_test
2024-11-25 18:04:52 +01:00
parent b76a292bde
commit 890ea7de93
1 changed files with 2 additions and 2 deletions
--- a/tests/quantization/ggml/test_ggml.py
+++ b/tests/quantization/ggml/test_ggml.py
@@ -623,8 +623,8 @@ class GgufIntegrationTests(unittest.TestCase):
            torch_dtype=torch.float16,
        )

-        text = tokenizer(self.example_text, return_tensors="pt").to(torch_device)
-        out = model.generate(**text, max_new_tokens=10)
+        text = tokenizer(self.example_text, return_tensors="pt")["input_ids"].to(torch_device)
+        out = model.generate(text, max_new_tokens=10)

        EXPECTED_TEXT = "Hello All,\nI am new to this forum."
        self.assertEqual(tokenizer.decode(out[0], skip_special_tokens=True), EXPECTED_TEXT)