fix sampler test

alibaba · Oct 9, 2024 · e4be207 · e4be207
1 parent 63f431a
commit e4be207
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tzrec/datasets/sampler_test.py b/tzrec/datasets/sampler_test.py
@@ -457,9 +457,9 @@ def _sampler_worker(pos_res, neg_res):
         p.join()
         if p.exitcode != 0:
             raise RuntimeError("worker failed.")
-        self.assertEqual(len(pos_res["int_a"]), 4 * 5)
-        self.assertEqual(len(pos_res["float_b"]), 4 * 5)
-        self.assertEqual(len(pos_res["str_c"]), 4 * 5)
+        self.assertEqual(len(pos_res["int_a"]), 4 * 4)
+        self.assertEqual(len(pos_res["float_b"]), 4 * 4)
+        self.assertEqual(len(pos_res["str_c"]), 4 * 4)
         self.assertEqual(len(neg_res["int_a"]), 4 * 15)
         self.assertEqual(len(neg_res["float_b"]), 4 * 15)
         self.assertEqual(len(neg_res["str_c"]), 4 * 15)