Merge pull request #302 from alexhernandezgarcia/fix_replay_bug

Fix permutation sampling in buffer select
alexhernandezgarcia · Apr 4, 2024 · ba8129a · ba8129a
2 parents 030f207 + 9274c37
commit ba8129a
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/gflownet/utils/buffer.py b/gflownet/utils/buffer.py
@@ -319,7 +319,12 @@ def select(
             samples = list(samples.values())
         if mode == "permutation":
             assert rng is not None
-            samples = [samples[idx] for idx in rng.permutation(n)]
+            indices = rng.choice(
+                len(samples),
+                size=n,
+                replace=False,
+            )
+            samples = [samples[idx] for idx in indices]
         elif mode == "weighted":
             if "rewards" in data_dict:
                 score = "rewards"