diff --git a/internlm/data/utils.py b/internlm/data/utils.py
index 3eee9d9..724fb9f 100644
--- a/internlm/data/utils.py
+++ b/internlm/data/utils.py
@@ -5,7 +5,7 @@ import torch
 
 from internlm.core.context import global_context as gpc
 
-DATASET_TYPE_IDS_MAP = {"en": 0, "cn": 1}
+DATASET_TYPE_IDS_MAP = {"en": 0, "cn": 1, "code": 2}
 
 
 def get_dataset_type_id(path):