diff --git a/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py b/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
index 34a037ab..47ec7072 100644
--- a/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
+++ b/server/text_generation_server/models/custom_modeling/flash_rw_modeling.py
@@ -101,11 +101,6 @@ class FlashRWAttention(torch.nn.Module):
     def __init__(
         self,
         config, prefix, weights,
-        # num_heads,
-        # num_heads_kv,
-        # hidden_size,
-        # bias,
-        # process_group=None,
         reduce=True,
     ):
         super().__init__()