Fix StoppingCriteriaSub parameters to be compatible with latest Trans…

…formers (#215)
intel · May 11, 2024 · a61d89f · a61d89f
1 parent cc1556d
commit a61d89f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llm_on_ray/inference/utils.py b/llm_on_ray/inference/utils.py
@@ -95,7 +95,7 @@ def __init__(self, stops=[], encounters=1):
         super().__init__()
         self.stops = stops
 
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs):
         for stop in self.stops:
             length = 1 if len(stop.size()) == 0 else stop.size()[0]
             if torch.all((stop == input_ids[0][-length:])).item():