From 0c7cd08255a4d5234a55905b4f037daea0f4ca62 Mon Sep 17 00:00:00 2001
From: BlinkDL <a@a.com>
Date: Fri, 24 Mar 2023 07:49:33 +0000
Subject: [PATCH] fix

---
 RWKV-v4neo/src/dataset.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/RWKV-v4neo/src/dataset.py b/RWKV-v4neo/src/dataset.py
index 71a8275..5261586 100644
--- a/RWKV-v4neo/src/dataset.py
+++ b/RWKV-v4neo/src/dataset.py
@@ -189,6 +189,9 @@ class MyDataset(Dataset):
                         i = np.random.randint(0, self.data_size - req_len)
                     else:
                         i = np.random.randint(0, self.data_size)
+                else:
+                    # cheat: pick a random spot in dataset
+                    i = np.random.randint(0, self.data_size - req_len)
 
                 if args.data_type == "binidx":
                     if args.my_pile_version == 1: