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

…addlePaddle#60616)
zhaohaixu · Jan 12, 2024 · 47d2bad · 47d2bad
1 parent e50f43e
commit 47d2bad
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