update tensor_parallel.py

server/text_generation_server/layers/tensor_parallel.py

@@ -65,7 +65,10 @@ class TensorParallelHead(SuperLayer):
         else:
             weight = weights.get_tensor(f"{prefix}.weight")
             should_gather = False
-
+        
+        if config.model_type == "baichuan":
+            weight = F.normalize(weight)
+        
         return TensorParallelHead(
             get_linear(weight, bias=None),
             process_group=weights.process_group,