Update train_model.py

train_model.py

@@ -67,9 +67,9 @@ def load_and_prepare_dataset(task, dataset_name, tokenizer, sequence_length):
     try:
         if '/' in dataset_name:
             dataset, config = dataset_name.split('/', 1)
-            dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
+            dataset = load_dataset("stanfordnlp/imdb",split='train')
         else:
-            dataset = load_dataset("Salesforce/wikitext", "wikitext-103-raw-v1", split='train')
+            dataset = load_dataset("stanfordnlp/imdb",split='train')
         
         logging.info("Dataset loaded successfully.")