From f6e33ddee63572d7571242298e579008886c12d3 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Thu, 12 Oct 2023 13:09:54 -0500 Subject: [PATCH] Fix tokens array dimensions in is_finished test The dimensions of the mock decoder output and the test input token ids did not match and were causing PyTorch to issue a warning/error in CI. --- tests/text/generate/test_continuation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/text/generate/test_continuation.py b/tests/text/generate/test_continuation.py index 78a3e34a8..944ae6df6 100644 --- a/tests/text/generate/test_continuation.py +++ b/tests/text/generate/test_continuation.py @@ -53,7 +53,7 @@ def test_continuation_stop_is_finished(): model = continuation(model, stop=["\n"]) - token_ids = torch.tensor([[2, 3]]) + token_ids = torch.tensor([[2, 3], [2, 3]]) result = model.is_finished(token_ids) assert torch.equal(result, torch.tensor([True, False]))