fix ReadTensor not reading the full contents of reader

PiperOrigin-RevId: 201040414

1 files changed, 1 insertions, 5 deletions

diff --git a/tensorflow/go/tensor.go b/tensorflow/go/tensor.go
index 2d25c04dc9..f3338f6595 100644
--- a/tensorflow/go/tensor.go
+++ b/tensorflow/go/tensor.go
@@ -131,13 +131,9 @@ func ReadTensor(dataType DataType, shape []int64, r io.Reader) (*Tensor, error)
 	}
 	runtime.SetFinalizer(t, (*Tensor).finalize)
 	raw := tensorData(t.c)
-	n, err := r.Read(raw)
-	if err != nil {
+	if _, err := io.ReadFull(r, raw); err != nil {
 		return nil, err
 	}
-	if uintptr(n) != nbytes {
-		return nil, fmt.Errorf("expected serialized tensor to be %v bytes, read %v", nbytes, n)
-	}
 	return t, nil
 }