correct (#13304)

2021-08-30 12:02:08 +02:00
parent 4046e66e40
commit 4362ee298a
1 changed files with 3 additions and 0 deletions
--- a/tests/test_modeling_rag.py
+++ b/tests/test_modeling_rag.py
@@ -988,6 +988,9 @@ class RagModelIntegrationTests(unittest.TestCase):
            torch_device
        )

+        if torch_device == "cuda":
+            rag_token.half()
+
        input_dict = tokenizer(
            self.test_data_questions,
            return_tensors="pt",