Place biased_logits on device

outlines-dev · Dec 8, 2023 · e0b3d87 · e0b3d87
1 parent fd5af1a
commit e0b3d87
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/outlines/generate/generator.py b/outlines/generate/generator.py
@@ -284,7 +284,7 @@ def bias_logits(
     A view of the original logits tensor where some values are masked.
 
     """
-    biased_logits = torch.empty(logits.shape)
+    biased_logits = torch.empty(logits.shape, device=logits.device)
     for i, ids in enumerate(ids_to_mask):
         mask = torch.full((logits.shape[-1],), -math.inf, device=logits.device)
         mask[ids] = 0