From 7f5087e6217458eec9335811be8499f9a368afa9 Mon Sep 17 00:00:00 2001
From: Siming Dai <908660116@qq.com>
Date: Wed, 3 Jul 2024 11:08:49 +0800
Subject: [PATCH] [Safetensors] Fix safetensors shape (#8702) (#8704)

* Update sequence_parallel for predict

* Do not save moe_group

* Fix safetensors reading
---
 paddlenlp/utils/safetensors.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/paddlenlp/utils/safetensors.py b/paddlenlp/utils/safetensors.py
index c273d0d973c2..54256023db7d 100644
--- a/paddlenlp/utils/safetensors.py
+++ b/paddlenlp/utils/safetensors.py
@@ -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:
@@ -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))