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
This commit is contained in:
James Keeling 2018-12-19 12:17:37 -08:00 committed by TensorFlower Gardener
parent f02f163f72
commit 5df47b8808

View File

@ -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
}