Hi @valhalla
Thanks for your suggestion. I modified the get_dataset function in run_language_modeling.py using datasets as explained here 38
dataset = load_dataset('text', data_files=file_path, split='train')
dataset = dataset.map(lambda ex: tokenizer(ex["text"], add_special_tokens=True,
                                               truncation=True, max_length=args.block_size,
                                               ), batched=True)
dataset.set_format(type='torch', columns=['input_ids'])

it runs without error on TPU but it is superslow. I used wikitext103 to test the code. Each training step takes less than 1 sec in the original code, but with datasets it takes more than 60 sec. Am I missing something or this is expected?