From 5df47b8808bf29d9b9ae9d7a56448dd70a5687be Mon Sep 17 00:00:00 2001 From: James Keeling Date: Wed, 19 Dec 2018 12:17:37 -0800 Subject: [PATCH] Fix memory leak in Go API. The code previously wrapped a C pointer without setting the enclosing TensorHandle's finalizer. I add a new method, newTensorHandleFromC, to do this. PiperOrigin-RevId: 226213607 --- tensorflow/go/tensor_handle.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tensorflow/go/tensor_handle.go b/tensorflow/go/tensor_handle.go index befc1c43ba1..3b06773dd16 100644 --- a/tensorflow/go/tensor_handle.go +++ b/tensorflow/go/tensor_handle.go @@ -59,6 +59,13 @@ func (th *TensorHandle) finalizer() { C.TFE_DeleteTensorHandle(th.c) } +// newTensorHandleFromC takes ownership of c and returns the owning TensorHandle. +func newTensorHandleFromC(c *C.TFE_TensorHandle) *TensorHandle { + th := &TensorHandle{c: c} + runtime.SetFinalizer(th, (*TensorHandle).finalizer) + return th +} + // DataType returns the TensorHandle's datatype. func (th *TensorHandle) DataType() DataType { return DataType(C.TFE_TensorHandleDataType(th.c)) @@ -150,5 +157,5 @@ func (th *TensorHandle) CopyToDevice(c *Context, deviceName string) (*TensorHand if err := status.Err(); err != nil { return nil, err } - return &TensorHandle{c: newTh}, nil + return newTensorHandleFromC(newTh), nil }