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:
parent
f02f163f72
commit
5df47b8808
@ -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
|
||||
}
|
||||
|
Loading…
Reference in New Issue
Block a user