Update test

neuml · Dec 3, 2024 · b05abd2 · b05abd2
1 parent c6cd565
commit b05abd2
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/test/python/testpipeline/testdata/testtextractor.py b/test/python/testpipeline/testdata/testtextractor.py
@@ -4,6 +4,8 @@
 
 import unittest
 
+import torch
+
 from txtai.pipeline import Textractor
 
 # pylint: disable = C0411
@@ -54,6 +56,9 @@ def testDocling(self):
         Test docling backend
         """
 
+        # Set default device to CPU
+        torch.set_default_device("cpu")
+
         textractor = Textractor(backend="docling")
 
         # Extract text and check for Markdown formatting