diff --git a/function.py b/function.py
index fca3839..0c8d3ec 100644
--- a/function.py
+++ b/function.py
@@ -33,21 +33,21 @@ def load_data(filepath, k_features=20):
                     except ValueError:
                         features.append(np.nan)  # 用 NaN 表示缺失值
                 label = int(line[-1])
-                data.append(features)
                 labels.append(label)
+                data.append(features)
 
     # 转换为 numpy 数组
     data = np.array(data)
     labels = np.array(labels)
 
-    # 使用SimpleImputer填充NaN
+    # 填充缺失值
     imputer = SimpleImputer(strategy='mean')
     data = imputer.fit_transform(data)
 
-    # 特征标准化，避免除零错误
+    # 特征标准化
     mean = data.mean(axis=0)
     std_dev = data.std(axis=0)
-    std_dev[std_dev == 0] = 1  # 如果标准差为0，将其设为1，避免除以零
+    std_dev[std_dev == 0] = 1  # 避免除以零
     data = (data - mean) / std_dev
 
     # 特征选择
@@ -116,10 +116,14 @@ def predict_logreg(inX, weights, threshold=0.5):
 
 # 测试代码
 if __name__ == "__main__":
-    # 加载数据并选择前20个最重要的特征
+    # 加载数据
     X_train, y_train = load_data("./data/horse-colic.data", k_features=20)
     X_test, y_test = load_data("./data/horse-colic.test", k_features=20)
 
+    # 将标签转换为二分类的0和1
+    y_train = np.where(y_train == 2, 0, 1)
+    y_test = np.where(y_test == 2, 0, 1)
+
     # 训练模型
     weights = train_logreg(X_train, y_train, initial_lr=0.005, iterations=5000, regularization=0.01, decay=0.0001)