Internal change

PiperOrigin-RevId: 373387638
tensorflow · May 12, 2021 · 1473751 · 1473751
1 parent 857c457
commit 1473751
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/documentation/_index.yaml b/documentation/_index.yaml
@@ -50,7 +50,7 @@ landing_page:
           <li><a href="migration">Migration guide from Neural Network to Decision Forests</a></li>
           <li><a href="https://github.com/tensorflow/decision-forests/issues">Issue tracker</a></li>
           <li><a href="https://github.com/tensorflow/decision-forests/documentation/known_issues.md">Known issues</a></li>
-          <li><a href="https://github.com/tensorflow/decision-forests/blob/master/CHANGELOG">Changelog</a></li>
+          <li><a href="https://github.com/tensorflow/decision-forests/blob/master/CHANGELOG.md">Changelog</a></li>
           <li><a href="https://discuss.tensorflow.org">Discuss on TensorFlow.Org</a></li>
           <li><a href="https://github.com/google/yggdrasil-decision-forests">Yggdrasil Decision Forest</a> (for advanced users)</li>
           </ul>

diff --git a/tensorflow_decision_forests/tensorflow/ops/inference/kernel.cc b/tensorflow_decision_forests/tensorflow/ops/inference/kernel.cc
@@ -1189,9 +1189,12 @@ class SimpleMLCreateModelResource : public OpKernel {
  public:
   explicit SimpleMLCreateModelResource(OpKernelConstruction* ctx)
       : OpKernel(ctx), model_handle_set_(false) {
-    OP_REQUIRES_OK(ctx, ctx->allocate_persistent(tensorflow::DT_RESOURCE,
-                                                 tensorflow::TensorShape({}),
-                                                 &model_handle_, nullptr));
+    // The model_handle_ object will outlive the constructor function.
+    // It will live until the resource is freed manually or by
+    // explicit ~Tensor() destructor.
+    OP_REQUIRES_OK(
+        ctx, ctx->allocate_temp(tensorflow::DT_RESOURCE,
+                                tensorflow::TensorShape({}), &model_handle_));
   }
 
   ~SimpleMLCreateModelResource() override {
@@ -1235,18 +1238,17 @@ class SimpleMLCreateModelResource : public OpKernel {
     tf::core::ScopedUnref unref_me(model);
 
     if (!model_handle_set_) {
-      auto h = model_handle_.AccessTensor(ctx)
-                   ->template scalar<tf::ResourceHandle>();
+      auto h = model_handle_.template scalar<tf::ResourceHandle>();
       h() = tf::MakeResourceHandle<YggdrasilModelResource>(
           ctx, cinfo_.container(), cinfo_.name());
     }
-    ctx->set_output(0, *model_handle_.AccessTensor(ctx));
+    ctx->set_output(0, model_handle_);
     model_handle_set_ = true;
   }
 
  private:
   tf::mutex mu_;
-  tf::PersistentTensor model_handle_ TF_GUARDED_BY(mu_);
+  tf::Tensor model_handle_ TF_GUARDED_BY(mu_);
   bool model_handle_set_ TF_GUARDED_BY(mu_);
   tf::ContainerInfo cinfo_;