Fix samplers_test

PiperOrigin-RevId: 428463103
google-deepmind · Feb 14, 2022 · 0840c98 · 0840c98
1 parent d2731dd
commit 0840c98
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/clrs/_src/samplers_test.py b/clrs/_src/samplers_test.py
@@ -26,9 +26,11 @@
 
 class SamplersTest(parameterized.TestCase):
 
-  @parameterized.parameters(*specs.CLRS_21_ALGS)
+  @parameterized.parameters(*specs.CLRS_30_ALGS)
   def test_sampler_determinism(self, name):
-    sampler, _ = samplers.clrs21_val(name)
+    num_samples = 3
+    num_nodes = 10
+    sampler, _ = samplers.build_sampler(name, num_samples, num_nodes)
 
     np.random.seed(47)  # Set seed
     feedback = sampler.next()