From 4e74dc86af32ccfabcf61b467257b17ea563861c Mon Sep 17 00:00:00 2001
From: "Quan (Andy) Gan" <coin2028@hotmail.com>
Date: Fri, 2 Jul 2021 13:37:06 +0800
Subject: [PATCH] fix curand (#3077)

Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
---
 src/random/random.cc | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/random/random.cc b/src/random/random.cc
index 975c6d162082..e6436536015c 100644
--- a/src/random/random.cc
+++ b/src/random/random.cc
@@ -24,16 +24,16 @@ DGL_REGISTER_GLOBAL("rng._CAPI_SetSeed")
 #pragma omp parallel for
     for (int i = 0; i < omp_get_max_threads(); ++i) {
       RandomEngine::ThreadLocal()->SetSeed(seed);
+    }
 #ifdef DGL_USE_CUDA
-      auto* thr_entry = CUDAThreadEntry::ThreadLocal();
-      if (!thr_entry->curand_gen) {
-        CURAND_CALL(curandCreateGenerator(&thr_entry->curand_gen, CURAND_RNG_PSEUDO_DEFAULT));
-      }
-      CURAND_CALL(curandSetPseudoRandomGeneratorSeed(
-          thr_entry->curand_gen,
-          static_cast<uint64_t>(seed + GetThreadId())));
-#endif  // DGL_USE_CUDA
+    auto* thr_entry = CUDAThreadEntry::ThreadLocal();
+    if (!thr_entry->curand_gen) {
+      CURAND_CALL(curandCreateGenerator(&thr_entry->curand_gen, CURAND_RNG_PSEUDO_DEFAULT));
     }
+    CURAND_CALL(curandSetPseudoRandomGeneratorSeed(
+        thr_entry->curand_gen,
+        static_cast<uint64_t>(seed)));
+#endif  // DGL_USE_CUDA
   });
 
 DGL_REGISTER_GLOBAL("rng._CAPI_Choice")