diff --git a/wit/lit_module.py b/wit/lit_module.py
index af2fc4d..5e3852c 100644
--- a/wit/lit_module.py
+++ b/wit/lit_module.py
@@ -63,11 +63,10 @@ class LitModule(pl.LightningModule):
         logits = logits.contiguous().view(-1, logits.size(-1))
         labels = batch["labels"][..., 1:]
         labels = labels.contiguous().view(-1)
-        label_mask = labels < self.vocab_size
+        label_mask = batch["mask"][..., 1:]
+        label_mask = label_mask.contiguous().view(-1)
         logits_m = logits[label_mask]
         labels_m = labels[label_mask]
-        # m = torch.max(logits, 1).indices.cpu().numpy()
-        # ll = labels.cpu().numpy()
         self.metric_accuracy.update(logits_m, labels_m)
         self.metric_loss.update(loss)
 
diff --git a/wit/meaning_dataset.py b/wit/meaning_dataset.py
index 416134e..7b49792 100644
--- a/wit/meaning_dataset.py
+++ b/wit/meaning_dataset.py
@@ -241,6 +241,8 @@ class MeaningDataset(Dataset):
         map = MeaningMap(size=end, vocab_size=vocab_size, max_subitem=max_subitem, use_cache=use_cache)
         np.random.seed(seed)
 
+        self.mask_level = None
+        self.mask_idx = None
         self.tree = []
         self.seq = []
         self.level = []
@@ -277,6 +279,10 @@ class MeaningDataset(Dataset):
     def len(self):
         return len(self.seq)
 
+    def set_mask(self, level=None, idx=None):
+        self.mask_level = level
+        self.mask_idx = idx
+
     def __getitem__(self, idx):
         output = {}
         data = torch.tensor(self.seq[idx]).long()
@@ -285,6 +291,10 @@ class MeaningDataset(Dataset):
         output["token_type_ids"] = torch.zeros(data.shape)
         output["tree"] = self.tree[idx]
         output["level"] = self.level[idx]
+        if self.mask_level is not None and self.mask_idx is not None:
+            output["mask"] = torch.tensor(self.get_seq_mask(idx, self.mask_level, self.mask_idx))
+        else:
+            output["mask"] = torch.ones(data.shape, dtype=torch.long)
         return output
 
     def get_batch(self, idx_list):  # must equal sequence length
@@ -296,6 +306,12 @@ class MeaningDataset(Dataset):
         output["token_type_ids"] = torch.zeros(data.shape)
         output["tree"] = [self.tree[i] for i in idx_list]
         output["level"] = [self.level[i] for i in idx_list]
+        if self.mask_level is not None and self.mask_idx is not None:
+            output["mask"] = torch.tensor(
+                np.stack([self.get_seq_mask(i, self.mask_level, self.mask_idx) for i in idx_list], axis=0)
+            )
+        else:
+            output["mask"] = torch.ones(data.shape, dtype=torch.long)
         return output
 
     def get_token(self, idx):  # must equal sequence length
diff --git a/wit/train.py b/wit/train.py
index 1260e92..780392e 100644
--- a/wit/train.py
+++ b/wit/train.py
@@ -34,9 +34,14 @@ hidden_size = 1024  # 128 1024 2048  32
 num_attention_heads = 16  # 8 8 16
 num_hidden_layers = 3  # 6 12 24  3
 
-name = "vocab_ratio_level_data_hidden_head_layer"
+mask_level = None
+mask_idx = None
+
+# name = "vocab_ratio_level_data_hidden_head_layer"
+name = "rank"
 ver = f"{vocab_size}" + "_" + f"{level_ratio}" + "_" + f"{level}" + "_" + f"{dataset_level}"
 ver = ver + "_" + f"{hidden_size}" + "_" + f"{num_attention_heads}" + "_" + f"{num_hidden_layers}"
+ver = ver + "_" + f"{mask_level}" + "_" + f"{mask_idx}"
 
 if __name__ == "__main__":
     torch.manual_seed(seed)
@@ -53,6 +58,7 @@ if __name__ == "__main__":
     start = vocab_size * (level_ratio**level)
     size = vocab_size * (level_ratio**dataset_level)
     raw_dataset = MeaningDataset(start, start + size, size, vocab_size, level_ratio)
+    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)