yikes there are even more random seeds!

ChiSym · Jun 18, 2024 · 86b6d4e · 86b6d4e
1 parent d0b71e6
commit 86b6d4e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/genparse/steer.py b/genparse/steer.py
@@ -29,9 +29,11 @@
 
 def set_seed(seed):
     random.seed(seed)
+    np.random.seed(seed)
     torch.manual_seed(seed)
     transformers.set_seed(seed)
-    np.random.seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
 
 
 # ____________________________________________________________________________________