style

HabanaAI · Dec 11, 2024 · 91c1dfe · 91c1dfe
1 parent 19d4e14
commit 91c1dfe
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/vllm/model_executor/layers/sampler.py b/vllm/model_executor/layers/sampler.py
@@ -405,10 +405,10 @@ def __init__(self, increment: int):
         self._increment = increment
 
     def __call__(self, logits: torch.Tensor, p: float, k: int):
-        if k==1 and not ApplyToppTopkScalar._handle_duplicates:
+        if k == 1 and not ApplyToppTopkScalar._handle_duplicates:
             new_logits = torch.full(logits.shape,
-                            -float("inf"),
-                            device=logits.device)
+                                    -float("inf"),
+                                    device=logits.device)
             vals, idx = torch.max(logits, keepdim=True, dim=1)
             new_logits.scatter_(1, idx, vals.to(new_logits.dtype))
             return new_logits