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