Skip to content

Commit

Permalink
[Safetensors] Fix safetensors shape (PaddlePaddle#8702)
Browse files Browse the repository at this point in the history
* Update sequence_parallel for predict

* Do not save moe_group

* Fix safetensors reading
  • Loading branch information
DesmonDay committed Jul 3, 2024
1 parent db99efd commit ad271a6
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions paddlenlp/utils/safetensors.py
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,7 @@ def __getitem__(self, index):
span = self.bits
for i, (start, stop, step) in enumerate(zip(out_start[::-1], out_stop[::-1], out_step[::-1])):
if len(indices) == 0:
if start == 0 and stop == self.shape[i]:
if start == 0 and stop == self.shape[::-1][i]:
pass
# We haven't started to slice yet, just increase the span
else:
Expand All @@ -194,7 +194,7 @@ def __getitem__(self, index):
newindices.append((old_start + offset, old_stop + offset))
indices = newindices
assert len(indices) == capacity, f"error {capacity} {len(indices)}"
span *= self.shape[-(i + 1)]
span *= self.shape[::-1][i]

if len(indices) == 0:
indices.append((0, self.nbytes))
Expand Down

0 comments on commit ad271a6

Please sign in to comment.