diff --git a/pytorch_mebeauty_dataset.py b/pytorch_mebeauty_dataset.py
index 6436fac..004b32f 100644
--- a/pytorch_mebeauty_dataset.py
+++ b/pytorch_mebeauty_dataset.py
@@ -88,9 +88,9 @@ def train_test_data(train_scores, test_scores,root_dir = '', train_augmentation
     train_data = MEBeauty(root_dir, train_scores,test_scores, train = True, transform = transform_train)
     trainloader = torch.utils.data.DataLoader(train_data, batch_size = batch, shuffle =True, num_workers = num_workers, pin_memory = pin_memory)
     test_data = MEBeauty(root_dir, train_scores,test_scores, train = False, transform = transform_test)
-    testloader = torch.utils.data.DataLoader(train_data, batch_size = batch, shuffle =True, num_workers = num_workers, pin_memory = pin_memory)
+    testloader = torch.utils.data.DataLoader(test_data, batch_size = batch, shuffle =True, num_workers = num_workers, pin_memory = pin_memory)
     
-    return trainloader, trainloader
+    return trainloader, testloader
   
 def plot_training(train_losses, valid_losses):
     
@@ -100,4 +100,4 @@ def plot_training(train_losses, valid_losses):
     plt.plot(train_losses, label="train_loss")
     plt.plot(valid_losses, label="valid_loss")
     plt.legend()
-   
\ No newline at end of file
+