fix: rare device mismatch

AnswerDotAI · Feb 23, 2024 · 4db672e · 4db672e
1 parent bfd0d57
commit 4db672e
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/ragatouille/models/colbert.py b/ragatouille/models/colbert.py
@@ -11,7 +11,6 @@
 from colbert import Indexer, IndexUpdater, Searcher, Trainer
 from colbert.infra import ColBERTConfig, Run, RunConfig
 from colbert.modeling.checkpoint import Checkpoint
-
 from ragatouille.models.base import LateInteractionModel
 
 # TODO: Move all bsize related calcs to `_set_bsize()`
@@ -751,7 +750,7 @@ def encode(
                         - encodings.shape[1],
                         encodings.shape[2],
                     )
-                ),
+                ).to(device=encodings.device),
             ],
             dim=1,
         )
@@ -765,7 +764,7 @@ def encode(
                         - doc_masks.shape[1],
                     ),
                     -float("inf"),
-                ),
+                ).to(device=encodings.device),
             ],
             dim=1,
         )