Create a new curand instead of reseeding. (#1089)

huggingface · Oct 14, 2023 · 9309cfc · 9309cfc
1 parent a193bf5
commit 9309cfc
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/candle-core/src/cuda_backend.rs b/candle-core/src/cuda_backend.rs
@@ -224,8 +224,10 @@ impl BackendDevice for CudaDevice {
     }
 
     fn set_seed(&self, seed: u64) -> Result<()> {
+        // We do not call set_seed but instead create a new curand object. This ensures that the
+        // state will be identical and the same random numbers will be generated.
         let mut curand = self.curand.lock().unwrap();
-        curand.0.set_seed(seed).w()?;
+        curand.0 = cudarc::curand::CudaRng::new(seed, self.device.clone()).w()?;
         Ok(())
     }
 

diff --git a/candle-core/src/device.rs b/candle-core/src/device.rs
@@ -128,6 +128,13 @@ impl Device {
         Ok(Self::Cuda(crate::CudaDevice::new(ordinal)?))
     }
 
+    pub fn set_seed(&self, seed: u64) -> Result<()> {
+        match self {
+            Self::Cpu => crate::cpu_backend::CpuDevice.set_seed(seed),
+            Self::Cuda(c) => c.set_seed(seed),
+        }
+    }
+
     pub fn same_device(&self, rhs: &Self) -> bool {
         match (self, rhs) {
             (Self::Cpu, Self::Cpu) => true,