diff --git a/server/text_generation_server/utils/layers.py b/server/text_generation_server/utils/layers.py
index 27adc775..8e36f654 100644
--- a/server/text_generation_server/utils/layers.py
+++ b/server/text_generation_server/utils/layers.py
@@ -355,8 +355,7 @@ class FastLinearROCm(nn.Module):
         weight = self.weight
         bias = self.bias
 
-        # TODO: fix for TP>=2, this only works for TP=1
-        if False and IS_ROCM_SYSTEM and inp.numel() // inp.size(-1) == 1:
+        if IS_ROCM_SYSTEM and inp.numel() // inp.size(-1) == 1:
             batched = False
 
             if inp.dim() == 3: