diff --git a/CAT/model/NCD.py b/CAT/model/NCD.py
index e4f6a83..1b946dd 100644
--- a/CAT/model/NCD.py
+++ b/CAT/model/NCD.py
@@ -13,13 +13,13 @@ class NCD(nn.Module):
     '''
     NeuralCDM
     '''
-    def __init__(self, student_n, exer_n, knowledge_n):
+    def __init__(self, student_n, exer_n, knowledge_n, prednet_len1=128, prednet_len2=64):
         self.knowledge_dim = knowledge_n
         self.exer_n = exer_n
         self.emb_num = student_n
         self.stu_dim = self.knowledge_dim
         self.prednet_input_len = self.knowledge_dim
-        self.prednet_len1, self.prednet_len2 = 512, 256  # changeable
+        self.prednet_len1, self.prednet_len2 = prednet_len1, prednet_len2  # changeable
 
         super(NCD, self).__init__()
 
@@ -96,7 +96,7 @@ def name(self):
         return 'Neural Cognitive Diagnosis'
 
     def init_model(self, data: Dataset):
-        self.model = NCD(data.num_students, data.num_questions, data.num_concepts)
+        self.model = NCD(data.num_students, data.num_questions, data.num_concepts, self.config['prednet_len1'], self.config['prednet_len2'])
     
     def train(self, train_data: TrainDataset):
         lr = self.config['learning_rate']
diff --git a/scripts/test.ipynb b/scripts/test.ipynb
index a6de80a..8f33565 100644
--- a/scripts/test.ipynb
+++ b/scripts/test.ipynb
@@ -107,8 +107,11 @@
     "    'learning_rate': 0.0025,\n",
     "    'batch_size': 2048,\n",
     "    'num_epochs': 8,\n",
-    "    'num_dim': 1,\n",
+    "    'num_dim': 1, # for IRT or MIRT\n",
     "    'device': 'cpu',\n",
+    "    # for NeuralCD\n",
+    "    'prednet_len1': 128,\n",
+    "    'prednet_len2': 64,\n",
     "}\n",
     "# fixed test length\n",
     "test_length = 5\n",
diff --git a/scripts/train.ipynb b/scripts/train.ipynb
index 4c902ba..c3a27cb 100644
--- a/scripts/train.ipynb
+++ b/scripts/train.ipynb
@@ -52,8 +52,11 @@
     "    'learning_rate': 0.002,\n",
     "    'batch_size': 2048,\n",
     "    'num_epochs': 10,\n",
-    "    'num_dim': 10,\n",
+    "    'num_dim': 10, # for IRT or MIRT\n",
     "    'device': 'cpu',\n",
+    "    # for NeuralCD\n",
+    "    'prednet_len1': 128,\n",
+    "    'prednet_len2': 64,\n",
     "}"
    ]
   },