mirror of
https://github.com/huggingface/text-generation-inference.git
synced 2025-09-10 11:54:52 +00:00
Update server/text_generation_server/utils/layers.py
This commit is contained in:
parent
63f03b4b7d
commit
6193512c4b
@ -286,7 +286,6 @@ class TensorParallelRowLinear(SuperLayer):
|
||||
class TensorParallelEmbedding(nn.Module):
|
||||
def __init__(self, prefix: str, weights, reduce=True):
|
||||
super().__init__()
|
||||
# weight = weights.get_sharded(f"{prefix}.weight", dim=0)
|
||||
weight = weights.get_partial_sharded(f"{prefix}.weight", dim=0)
|
||||
num_embeddings = weights.get_shape(f"{prefix}.weight")[0]
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user