Update train_mlp.py

train_mlp.py

@@ -105,7 +105,7 @@ def main():
 
     # Split the dataset into train and validation sets
     train_dataset = dataset['train']
-    val_dataset = dataset['validation']  # Assuming 'validation' is the correct key
+    val_dataset = dataset['valid']  # Assuming 'validation' is the correct key
 
     # Determine the number of classes
     num_classes = len(set(train_dataset['label']))