diff --git a/tensorflow/c/eager/BUILD b/tensorflow/c/eager/BUILD index 566fef023ea..224b36a170c 100644 --- a/tensorflow/c/eager/BUILD +++ b/tensorflow/c/eager/BUILD @@ -360,6 +360,7 @@ filegroup( exclude = [ "c_api_experimental.cc", "*test*", + "*dlpack*", ], ), visibility = ["//visibility:public"], diff --git a/tensorflow/c/eager/dlpack.cc b/tensorflow/c/eager/dlpack.cc index f78df720511..95213948859 100644 --- a/tensorflow/c/eager/dlpack.cc +++ b/tensorflow/c/eager/dlpack.cc @@ -59,6 +59,8 @@ void DLManagedTensorDeleter(DLManagedTensor* arg) { TFDLManagedTensorCtx* owner = static_cast(arg->manager_ctx); owner->handle->Unref(); + delete owner->handle; + delete owner->tensor.dl_tensor.shape; delete owner; } diff --git a/tensorflow/c/eager/dlpack.h b/tensorflow/c/eager/dlpack.h index 43c205e30c4..35dfb682114 100644 --- a/tensorflow/c/eager/dlpack.h +++ b/tensorflow/c/eager/dlpack.h @@ -24,11 +24,11 @@ namespace tensorflow { const char* const kDlTensorCapsuleName = "dltensor"; -void* TFE_HandleToDLPack(TFE_TensorHandle* h, TF_Status* status); +TF_CAPI_EXPORT extern void* TFE_HandleToDLPack(TFE_TensorHandle* h, TF_Status* status); -TFE_TensorHandle* TFE_HandleFromDLPack(void* dlm, TF_Status* status); +TF_CAPI_EXPORT extern TFE_TensorHandle* TFE_HandleFromDLPack(void* dlm, TF_Status* status); -void TFE_CallDLManagedTensorDeleter(void* dlm_ptr); +TF_CAPI_EXPORT extern void TFE_CallDLManagedTensorDeleter(void* dlm_ptr); } // namespace tensorflow