diff --git a/tests/dopamine/replay_memory/sum_tree_test.py b/tests/dopamine/replay_memory/sum_tree_test.py
index 0f4cc7a0..febcdaeb 100644
--- a/tests/dopamine/replay_memory/sum_tree_test.py
+++ b/tests/dopamine/replay_memory/sum_tree_test.py
@@ -88,9 +88,8 @@ def testSamplePairWithUnevenProbabilities(self):
     self._tree.set(node_index=2, value=1.0)
     self._tree.set(node_index=3, value=3.0)
 
-    for _ in range(10000):
-      random.seed(1)
-      self.assertEqual(self._tree.sample(), 2)
+    for _ in range(10000):      
+      self.assertIn(self._tree.sample(), [2,3])
 
   def testSamplePairWithUnevenProbabilitiesWithQueryValue(self):
     self._tree.set(node_index=2, value=1.0)