modified: datasets.py

autonlab · Mar 29, 2022 · e46028a · e46028a
1 parent c190832
commit e46028a
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/auton_survival/datasets.py b/auton_survival/datasets.py
@@ -38,7 +38,9 @@
 
 import torchvision
 
-def increase_censoring(e, t, p):
+def increase_censoring(e, t, p, random_seed=0):
+
+  np.random.seed(random_seed)
 
   uncens = np.where(e == 1)[0]
   mask = np.random.choice([False, True], len(uncens), p=[1-p, p])