From b97307e6316d6ebd1b8b5343bde02769e3052c99 Mon Sep 17 00:00:00 2001
From: zhaohx <zhaohx@tecorigin.com>
Date: Mon, 8 Jan 2024 16:01:18 +0800
Subject: [PATCH] This PR enable offset of generator for custom device.

---
 paddle/phi/core/generator.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/core/generator.cc b/paddle/phi/core/generator.cc
index ef5d3e7a646211..4aa3568ea6ae9e 100644
--- 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;