[cherry-pick] This PR enable offset of generator for custom device. (#…

…60616) (#60772)
PaddlePaddle · Jan 23, 2024 · 0f732a5 · 0f732a5
1 parent 1aa5f4b
commit 0f732a5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paddle/phi/core/generator.cc b/paddle/phi/core/generator.cc
@@ -278,7 +278,8 @@ uint64_t Generator::Random64() {
 
 std::pair<uint64_t, uint64_t> Generator::IncrementOffset(
     uint64_t increment_offset) {
-#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP) || \
+    defined(PADDLE_WITH_CUSTOM_DEVICE)
   std::lock_guard<std::mutex> lock(this->mu_);
   uint64_t cur_offset = this->state_.thread_offset;
   VLOG(10) << "cur_offset = " << cur_offset