We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 7d8a417 commit d73a412Copy full SHA for d73a412
examples/colpali/exact.py
@@ -1,4 +1,5 @@
1
from colpali_engine.models import ColQwen2, ColQwen2Processor
2
+from colpali_engine.utils.torch_utils import get_torch_device
3
from datasets import load_dataset
4
from pgvector.psycopg import register_vector, Bit
5
import psycopg
@@ -30,7 +31,7 @@
30
31
""")
32
33
-device = 'mps' if torch.backends.mps.is_available() else 'cpu'
34
+device = get_torch_device('auto')
35
model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval()
36
processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0')
37
0 commit comments