From 3959bc3f93f6b2e00cc47d70ab7d0a7e20bc3349 Mon Sep 17 00:00:00 2001
From: sugarme <thangtran@hotmail.com>
Date: Tue, 25 Jul 2023 13:58:19 +1000
Subject: [PATCH] fixed read npy failed with single element tensor with zero
 shape

---
 ts/npy.go | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/ts/npy.go b/ts/npy.go
index e621f84..9f8f2e8 100644
--- a/ts/npy.go
+++ b/ts/npy.go
@@ -306,6 +306,11 @@ func ReadNpy(filepath string) (*Tensor, error) {
 		return nil, err
 	}
 
+	// NOTE(TT.). case tensor 1 element with shape = []
+	if len(data) > 0 && len(header.shape) == 0 {
+		header.shape = []int64{1}
+	}
+
 	return OfDataSize(data, header.shape, header.descr)
 }
 
@@ -348,6 +353,11 @@ func ReadNpz(filePath string) ([]NamedTensor, error) {
 			return nil, err
 		}
 
+		// NOTE(TT.). case tensor 1 element with shape = []
+		if len(data) > 0 && len(header.shape) == 0 {
+			header.shape = []int64{1}
+		}
+
 		tensor, err := OfDataSize(data, header.shape, header.descr)
 		if err != nil {
 			return nil, err