deve

app.py

@@ -207,7 +207,7 @@ def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size
             del dataset['validation']
             tokenized_first_third = dataset.map(tokenize_function, batched=True,)
             
-            tokenized_first_third.save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
+            tokenized_first_third['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')                           
             print('DONE')
             return 'RUN AGAIN TO LOAD REST OF DATA'
         dataset = load_dataset(dataset_name.strip())