fix #59

2020-05-29 11:49:47 +08:00 · 2020-05-29 11:49:47 +08:00 · be9ce44290
commit be9ce44290
parent d2b2fa87c0
1 changed files with 1 additions and 1 deletions
--- a/tianshou/data/collector.py
+++ b/tianshou/data/collector.py
@ -402,7 +402,7 @@ class Collector(object):
                lens = [len(b) for b in self.buffer]
                total = sum(lens)
                batch_index = np.random.choice(
-                    total, batch_size, p=np.array(lens) / total)
+                    len(self.buffer), batch_size, p=np.array(lens) / total)
            else:
                batch_index = np.array([])
            batch_data = Batch()