convert: fix shape

9ea492f1 · Michael Yang · ac64cd4e · 9ea492f1
Commit 9ea492f1 authored Mar 10, 2024 by Michael Yang
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 3 deletions

convert/convert.go convert/convert.go +3 -3

No files found.
--- a/convert/convert.go
+++ b/convert/convert.go
@@ -103,9 +103,9 @@ func ReadSafeTensors(fn string, offset uint64) ([]llm.Tensor, uint64, error) {
 			return []llm.Tensor{}, 0, err
 		}

-		shape := [4]uint64{1, 1, 1, 1}
-		for cnt, s := range data.Shape {
-			shape[cnt] = uint64(s)
+		shape := []uint64{0, 0, 0, 0}
+		for i := range data.Shape {
+			shape[i] = uint64(data.Shape[i])
 		}

 		t := llm.Tensor{