This PR enable offset of generator for custom device. (#60616)

PaddlePaddle · Jan 10, 2024 · da5399a · da5399a
1 parent b622e96
commit da5399a
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/paddle/phi/core/generator.cc b/paddle/phi/core/generator.cc
@@ -271,7 +271,8 @@ uint64_t Generator::Random64() {
 }
 
 std::pair<uint64_t, uint64_t> Generator::IncrementOffset(uint64_t increment) {
-#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(mu_);
   uint64_t offset = state().offset;
   state().offset = offset + increment;