correct (#13304)

4362ee29 · Patrick von Platen · GitHub · 4046e66e · 4362ee29
Unverified Commit 4362ee29 authored Aug 30, 2021 by Patrick von Platen Committed by GitHub Aug 30, 2021
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 0 deletions

tests/test_modeling_rag.py tests/test_modeling_rag.py +3 -0

No files found.
--- 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",