@@ -1189,9 +1189,12 @@ class SimpleMLCreateModelResource : public OpKernel {
1189
1189
public:
1190
1190
explicit SimpleMLCreateModelResource (OpKernelConstruction* ctx)
1191
1191
: 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_));
1195
1198
}
1196
1199
1197
1200
~SimpleMLCreateModelResource () override {
@@ -1235,18 +1238,17 @@ class SimpleMLCreateModelResource : public OpKernel {
1235
1238
tf::core::ScopedUnref unref_me (model);
1236
1239
1237
1240
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>();
1240
1242
h () = tf::MakeResourceHandle<YggdrasilModelResource>(
1241
1243
ctx, cinfo_.container (), cinfo_.name ());
1242
1244
}
1243
- ctx->set_output (0 , * model_handle_. AccessTensor (ctx) );
1245
+ ctx->set_output (0 , model_handle_);
1244
1246
model_handle_set_ = true ;
1245
1247
}
1246
1248
1247
1249
private:
1248
1250
tf::mutex mu_;
1249
- tf::PersistentTensor model_handle_ TF_GUARDED_BY (mu_);
1251
+ tf::Tensor model_handle_ TF_GUARDED_BY (mu_);
1250
1252
bool model_handle_set_ TF_GUARDED_BY (mu_);
1251
1253
tf::ContainerInfo cinfo_;
1252
1254
0 commit comments