Skip to content

Commit d73a412

Browse files
committed
Updated ColPali example to use get_torch_device [skip ci]
1 parent 7d8a417 commit d73a412

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

examples/colpali/exact.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from colpali_engine.models import ColQwen2, ColQwen2Processor
2+
from colpali_engine.utils.torch_utils import get_torch_device
23
from datasets import load_dataset
34
from pgvector.psycopg import register_vector, Bit
45
import psycopg
@@ -30,7 +31,7 @@
3031
""")
3132

3233

33-
device = 'mps' if torch.backends.mps.is_available() else 'cpu'
34+
device = get_torch_device('auto')
3435
model = ColQwen2.from_pretrained('vidore/colqwen2-v1.0', torch_dtype=torch.bfloat16, device_map=device).eval()
3536
processor = ColQwen2Processor.from_pretrained('vidore/colqwen2-v1.0')
3637

0 commit comments

Comments
 (0)