Update train_model.py

train_model.py

@@ -12,7 +12,7 @@ from transformers import (
     DataCollatorForLanguageModeling,
     DataCollatorWithPadding,
 )
-from datasets import load_dataset, Dataset
+from datasets import load_dataset
 import torch
 import os
 from huggingface_hub import login, HfApi, HfFolder