diff --git a/wit/train.py b/wit/train.py
index 284290c..0202ad5 100644
--- a/wit/train.py
+++ b/wit/train.py
@@ -1,8 +1,3 @@
-import argparse
-from functools import partial
-from itertools import chain
-from typing import Dict, Tuple
-
 import pytorch_lightning as pl
 import torch
 
@@ -24,23 +19,24 @@ max_epochs = 1000
 strategy = "auto"
 resume_from_ckpt_path = None
 seed = 42
+dataloader_works = 2
 
 vocab_size = 256
-level_ratio = 6
-level = 4
+level_ratio = 5
+level = 5
 dataset_level = 1.5
 min_subitem = 2
 
-hidden_size = 1024  # 128 1024 2048  32
+hidden_size = 128  # 128 1024 2048  32
 num_attention_heads = 16  # 8 8 16
 num_hidden_layers = 6  # 6 12 24  3
 
-mask_level = [0]
-mask_idx = [-1]
+mask_level = [0, 1]
+mask_idx = [0, -1]
 
 # name = "vocab_ratio_level_data_hidden_head_layer"
 # name = "mask_level_idx"
-name = "small"
+name = "hard"
 
 ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{min_subitem}" + "_" + f"{dataset_level}"
 ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
@@ -63,8 +59,8 @@ if __name__ == "__main__":
     raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio, min_subitem)
     raw_dataset.set_mask(mask_level, mask_idx)
     train_dataset, val_dataset = raw_dataset.split(0.9)
-    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size)
-    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size)
+    train_dataloader = BatchGroupMeaningDataloader(train_dataset, train_batch_size).dataloader(dataloader_works)
+    val_dataloader = BatchGroupMeaningDataloader(val_dataset, val_batch_size).dataloader(dataloader_works)
 
     # for i in range(len(train_dataloader)):
     #     print(train_dataloader.print_mapping(i))