diff --git a/train.py b/train.py
index 92de58f..c11f252 100644
--- a/train.py
+++ b/train.py
@@ -27,7 +27,7 @@ use_cosine_similarity = config['use_cosine_similarity']
 
 data_augment = get_data_transform_opes(s=config['s'], crop_size=96)
 
-train_dataset = datasets.STL10('./data', split='train', download=True, transform=DataTransform(data_augment))
+train_dataset = datasets.STL10('./data', split='train+unlabeled', download=True, transform=DataTransform(data_augment))
 
 train_loader, valid_loader = get_train_validation_data_loaders(train_dataset, config)