Update logit_processors/logit_.py

logit_processors/logit_.py

@@ -44,7 +44,7 @@ class ThinkLogitsProcessor:
                 # If repetation detected, force </think>
                 if self.find_repeated_ngrams(past_token_ids, n=CHUNK_SIZE):
                     # Set all other logits to -inf except for </think>
-                    logits = torch.full_like(logits, torch.finfo(torch.bfloat16).min)
+                    logits = torch.full_like(logits, torch.finfo(logits.dtype).min)
                     logits[self.think_end_token] = 1.0
             else:
                 # ratio