experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix memory leak in Go API.

Browse Source 
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
This commit is contained in:
James Keeling 2018-12-19 12:17:37 -08:00 committed by TensorFlower Gardener
parent f02f163f72
commit 5df47b8808
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tensorflow/go/tensor_handle.go
View File

@ -59,6 +59,13 @@ func (th *TensorHandle) finalizer() {
C.TFE_DeleteTensorHandle(th.c) 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. // DataType returns the TensorHandle's datatype.
func (th *TensorHandle) DataType() DataType { func (th *TensorHandle) DataType() DataType {
return DataType(C.TFE_TensorHandleDataType(th.c)) 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 { if err := status.Err(); err != nil {
return nil, err return nil, err
} }
return &TensorHandle{c: newTh}, nil return newTensorHandleFromC(newTh), nil
} }