Browse Source

convert: fix shape

Michael Yang 1 year ago
parent
commit
9ea492f1ce
1 changed files with 3 additions and 3 deletions
  1. 3 3
      convert/convert.go

+ 3 - 3
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{