diff --git a/chapters/en/chapter7/3.mdx b/chapters/en/chapter7/3.mdx
index bcaab1b45..30e4161e1 100644
--- a/chapters/en/chapter7/3.mdx
+++ b/chapters/en/chapter7/3.mdx
@@ -523,6 +523,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return default_data_collator(features)
 ```
@@ -563,6 +564,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return tf_default_data_collator(features)
 ```
diff --git a/chapters/fr/chapter7/3.mdx b/chapters/fr/chapter7/3.mdx
index 664624b5b..4ca41a4b5 100644
--- a/chapters/fr/chapter7/3.mdx
+++ b/chapters/fr/chapter7/3.mdx
@@ -524,6 +524,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return default_data_collator(features)
 ```
@@ -564,6 +565,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return tf_default_data_collator(features)
 ```
diff --git a/chapters/ja/chapter7/3.mdx b/chapters/ja/chapter7/3.mdx
index 740c3f9bb..afdb30047 100644
--- a/chapters/ja/chapter7/3.mdx
+++ b/chapters/ja/chapter7/3.mdx
@@ -535,6 +535,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return default_data_collator(features)
 ```
@@ -575,6 +576,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return tf_default_data_collator(features)
 ```
diff --git a/chapters/vi/chapter7/3.mdx b/chapters/vi/chapter7/3.mdx
index 9a3048af6..e354003cf 100644
--- a/chapters/vi/chapter7/3.mdx
+++ b/chapters/vi/chapter7/3.mdx
@@ -523,6 +523,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return default_data_collator(features)
 ```
@@ -563,6 +564,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return tf_default_data_collator(features)
 ```
diff --git a/chapters/zh-CN/chapter7/3.mdx b/chapters/zh-CN/chapter7/3.mdx
index ced0bc325..d95d76313 100644
--- a/chapters/zh-CN/chapter7/3.mdx
+++ b/chapters/zh-CN/chapter7/3.mdx
@@ -524,6 +524,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return default_data_collator(features)
 ```
@@ -564,6 +565,7 @@ def whole_word_masking_data_collator(features):
             for idx in mapping[word_id]:
                 new_labels[idx] = labels[idx]
                 input_ids[idx] = tokenizer.mask_token_id
+        feature["labels"] = new_labels
 
     return tf_default_data_collator(features)
 ```