diff --git a/server/text_generation_server/utils/logits_process.py b/server/text_generation_server/utils/logits_process.py
index 47f3a33c..9a738190 100644
--- a/server/text_generation_server/utils/logits_process.py
+++ b/server/text_generation_server/utils/logits_process.py
@@ -197,7 +197,6 @@ class HeterogeneousTopKLogitsWarper(LogitsWarper):
     def __init__(
         self,
         top_k: List[int],
-        dtype: torch.dtype,
         device: torch.device,
         filter_value: float = -math.inf,
         min_tokens_to_keep: int = 1,
diff --git a/server/text_generation_server/utils/tokens.py b/server/text_generation_server/utils/tokens.py
index c4e45b72..118129ae 100644
--- a/server/text_generation_server/utils/tokens.py
+++ b/server/text_generation_server/utils/tokens.py
@@ -187,7 +187,7 @@ class HeterogeneousNextTokenChooser:
 
         if any([x != 0 for x in top_k]):
             do_sample = [sample or x != 0 for x, sample in zip(top_k, do_sample)]
-            warpers.append(HeterogeneousTopKLogitsWarper(top_k, dtype, device))
+            warpers.append(HeterogeneousTopKLogitsWarper(top_k, device))
 
         if any([x < 1.0 for x in top_p]):
             do_sample = [sample or x < 1.0 for x, sample in zip(top_p, do_sample)]