Skip to content

Commit 1473751

Browse files
committed
Internal change
PiperOrigin-RevId: 373387638
1 parent 857c457 commit 1473751

File tree

2 files changed

+10
-8
lines changed

2 files changed

+10
-8
lines changed

documentation/_index.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,7 @@ landing_page:
5050
<li><a href="migration">Migration guide from Neural Network to Decision Forests</a></li>
5151
<li><a href="https://github.com/tensorflow/decision-forests/issues">Issue tracker</a></li>
5252
<li><a href="https://github.com/tensorflow/decision-forests/documentation/known_issues.md">Known issues</a></li>
53-
<li><a href="https://github.com/tensorflow/decision-forests/blob/master/CHANGELOG">Changelog</a></li>
53+
<li><a href="https://github.com/tensorflow/decision-forests/blob/master/CHANGELOG.md">Changelog</a></li>
5454
<li><a href="https://discuss.tensorflow.org">Discuss on TensorFlow.Org</a></li>
5555
<li><a href="https://github.com/google/yggdrasil-decision-forests">Yggdrasil Decision Forest</a> (for advanced users)</li>
5656
</ul>

tensorflow_decision_forests/tensorflow/ops/inference/kernel.cc

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1189,9 +1189,12 @@ class SimpleMLCreateModelResource : public OpKernel {
11891189
public:
11901190
explicit SimpleMLCreateModelResource(OpKernelConstruction* ctx)
11911191
: OpKernel(ctx), model_handle_set_(false) {
1192-
OP_REQUIRES_OK(ctx, ctx->allocate_persistent(tensorflow::DT_RESOURCE,
1193-
tensorflow::TensorShape({}),
1194-
&model_handle_, nullptr));
1192+
// The model_handle_ object will outlive the constructor function.
1193+
// It will live until the resource is freed manually or by
1194+
// explicit ~Tensor() destructor.
1195+
OP_REQUIRES_OK(
1196+
ctx, ctx->allocate_temp(tensorflow::DT_RESOURCE,
1197+
tensorflow::TensorShape({}), &model_handle_));
11951198
}
11961199

11971200
~SimpleMLCreateModelResource() override {
@@ -1235,18 +1238,17 @@ class SimpleMLCreateModelResource : public OpKernel {
12351238
tf::core::ScopedUnref unref_me(model);
12361239

12371240
if (!model_handle_set_) {
1238-
auto h = model_handle_.AccessTensor(ctx)
1239-
->template scalar<tf::ResourceHandle>();
1241+
auto h = model_handle_.template scalar<tf::ResourceHandle>();
12401242
h() = tf::MakeResourceHandle<YggdrasilModelResource>(
12411243
ctx, cinfo_.container(), cinfo_.name());
12421244
}
1243-
ctx->set_output(0, *model_handle_.AccessTensor(ctx));
1245+
ctx->set_output(0, model_handle_);
12441246
model_handle_set_ = true;
12451247
}
12461248

12471249
private:
12481250
tf::mutex mu_;
1249-
tf::PersistentTensor model_handle_ TF_GUARDED_BY(mu_);
1251+
tf::Tensor model_handle_ TF_GUARDED_BY(mu_);
12501252
bool model_handle_set_ TF_GUARDED_BY(mu_);
12511253
tf::ContainerInfo cinfo_;
12521254

0 commit comments

Comments
 (0)