From bd46b2d9cad2b98b351b1b6d3ee5bfcfb6d2fd48 Mon Sep 17 00:00:00 2001
From: Artemiy Ryabinkov <getmylag@gmail.com>
Date: Tue, 19 Jan 2021 21:53:46 +0200
Subject: [PATCH] Go: Deallocate large TF_TString on Tensor finalization

---
 tensorflow/go/tensor.go | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/tensorflow/go/tensor.go b/tensorflow/go/tensor.go
index df5b34cec89..cfb389d472d 100644
--- a/tensorflow/go/tensor.go
+++ b/tensorflow/go/tensor.go
@@ -95,8 +95,17 @@ func NewTensor(value interface{}) (*Tensor, error) {
 		c:     C.TF_AllocateTensor(C.TF_DataType(dataType), shapePtr, C.int(len(shape)), C.size_t(nbytes)),
 		shape: shape,
 	}
-	runtime.SetFinalizer(t, (*Tensor).finalize)
+
 	raw := tensorData(t.c)
+
+	runtime.SetFinalizer(t, func(t *Tensor) {
+		if dataType == String {
+			t.clearTStrings(raw, nflattened)
+		}
+
+		t.finalize()
+	})
+
 	buf := bytes.NewBuffer(raw[:0:len(raw)])
 
 	if isAllArray(val.Type()) {
@@ -206,6 +215,14 @@ func newTensorFromC(c *C.TF_Tensor) *Tensor {
 	return t
 }
 
+func (t *Tensor) clearTStrings(raw []byte, n int64) {
+	tstrs := (*(*[]C.TF_TString)(unsafe.Pointer(&raw)))[:n]
+
+	for _, tstr := range tstrs {
+		C.TF_TString_Dealloc(&tstr)
+	}
+}
+
 func (t *Tensor) finalize() { C.TF_DeleteTensor(t.c) }
 
 // DataType returns the scalar datatype of the Tensor.