diff --git a/tensorflow/core/kernels/data/random_seed_ops.cc b/tensorflow/core/kernels/data/random_seed_ops.cc index 3b66f9e2388..7f1c2fa4238 100644 --- a/tensorflow/core/kernels/data/random_seed_ops.cc +++ b/tensorflow/core/kernels/data/random_seed_ops.cc @@ -72,6 +72,7 @@ void AnonymousSeedGeneratorHandleOp::Compute(OpKernelContext* ctx) { int64 seed2; OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kSeed2, &seed2)); // Seeds will be consumed by `CreateResource`, which is called via `Compute`. + mutex_lock l(mu_); seeds_ = absl::make_unique(seed, seed2); OP_REQUIRES_OK(ctx, ParseScalarArgument(ctx, kReshuffle, &reshuffle_)); AnonymousResourceOp::Compute(ctx); @@ -82,7 +83,8 @@ std::string AnonymousSeedGeneratorHandleOp::name() { return kSeedGenerator; } Status AnonymousSeedGeneratorHandleOp::CreateResource( OpKernelContext* ctx, std::unique_ptr flib_def, std::unique_ptr pflr, - FunctionLibraryRuntime* lib, SeedGeneratorManager** manager) { + FunctionLibraryRuntime* lib, SeedGeneratorManager** manager) + TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) { if (reshuffle_) { *manager = new SeedGeneratorManager(new RandomSeedGenerator(*seeds_)); } else { diff --git a/tensorflow/core/kernels/data/random_seed_ops.h b/tensorflow/core/kernels/data/random_seed_ops.h index e013cc055bd..68c7e88426f 100644 --- a/tensorflow/core/kernels/data/random_seed_ops.h +++ b/tensorflow/core/kernels/data/random_seed_ops.h @@ -142,7 +142,8 @@ class AnonymousSeedGeneratorHandleOp FunctionLibraryRuntime* lib, SeedGeneratorManager** manager) override; - std::unique_ptr seeds_ = nullptr; + mutex mu_; + std::unique_ptr seeds_ TF_GUARDED_BY(mu_); bool reshuffle_; };