diff --git a/gpt_2_simple/src/load_dataset.py b/gpt_2_simple/src/load_dataset.py
index 5694502..be264f2 100644
--- a/gpt_2_simple/src/load_dataset.py
+++ b/gpt_2_simple/src/load_dataset.py
@@ -26,7 +26,7 @@ def load_dataset(enc, path, combine):
     for path in tqdm.tqdm(paths):
         if path.endswith('.npz'):
             # Pre-encoded
-            with np.load(path) as npz:
+            with np.load(path,allow_pickle=True) as npz:
                 for item in npz.files:
                     token_chunks.append(npz[item])
         elif path.endswith('.csv'):