diff --git a/tokenise+train.py b/tokenise+train.py
index 35b227f..2d5ddfe 100644
--- a/tokenise+train.py
+++ b/tokenise+train.py
@@ -15,7 +15,7 @@ def suffix(bs: int, ns: int, vs: int) -> str:
 def train(filepath: str, ouputdir: Path, blocksize: int, vocabsize: int, num_steps: int, gpu: bool = False) -> str:
 
 	# from aitextgen.TokenDataset import TokenDataset
-	# from transformers import GPT2Config
+	from transformers import GPT2Config
 	from aitextgen.utils import build_gpt2_config
 	from aitextgen import aitextgen