remove unused cache

HabanaAI · Sep 23, 2024 · c04af23 · c04af23
1 parent 6d57c18
commit c04af23
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/vllm/model_executor/guided_decoding/outlines_logits_processors.py b/vllm/model_executor/guided_decoding/outlines_logits_processors.py
@@ -21,7 +21,6 @@
 from typing import Callable, DefaultDict, Dict, List, Union
 
 import torch
-from outlines.caching import cache
 from outlines.fsm.guide import CFGGuide, Generate, Guide, RegexGuide, Write
 from outlines.fsm.json_schema import build_regex_from_schema
 from pydantic import BaseModel
@@ -59,9 +58,11 @@ def __call__(self, input_ids: List[int],
             raise TypeError(
                 f"Unsupported instruction type {type(instruction)}")
 
-        mask = torch.ones((scores.shape[-1], ), device=scores.device, dtype=torch.bool)
+        mask = torch.full((scores.shape[-1], ),
+                          -math.inf,
+                          device=scores.device)
         mask[allowed_tokens] = 0
-        scores.masked_fill_(mask, -math.inf)
+        scores = scores.add(mask)
         return scores