Create mask on device

dottxt-ai · Jul 24, 2023 · a9cd45b · a9cd45b
1 parent cfeb812
commit a9cd45b
Showing 1 changed file with 1 addition and 3 deletions.
diff --git a/outlines/text/generate/regex.py b/outlines/text/generate/regex.py
@@ -128,9 +128,7 @@ def create_proposal(
         masks = []
         for pstate in self.pstates:
             mask = torch.full(
-                (len(self.model.tokenizer.vocabulary),),
-                -math.inf,
-                device=self.model.device,
+                (len(self.model.tokenizer.vocabulary),), -math.inf, device=self.device
             )
 
             if pstate[1] > -1: