Skip to content

Commit d5556db

Browse files
committed
TP: Fix sharing 0-dimensional tensors
1 parent a29d97a commit d5556db

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

exllamav3/model/model_tp_shared.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,8 @@ def send(self, tensor: torch.Tensor | None) -> dict:
6868
self.next_offset += nbytes_align
6969

7070
# Copy to shared buffer
71-
t_cpu = tensor.cpu().contiguous()
71+
tensor_d = tensor.view((1,)) if len(tensor.shape) == 0 else tensor
72+
t_cpu = tensor_d.cpu().contiguous()
7273
src = t_cpu.view(torch.uint8).numpy().view(np.uint8).ravel()
7374
dst = np.ndarray((nbytes,), dtype = np.uint8, buffer = self.shm.buf, offset = offset)
7475
np.copyto(dst, src, casting = "no")

0 commit comments

Comments
 (0)