diff --git a/train.py b/train.py
index 83e930c..57d6011 100644
--- a/train.py
+++ b/train.py
@@ -31,7 +31,7 @@ TRAIN_CONFIG = {
         "weight_decay": 0.001,     # 权重衰减
         "label_smoothing": 0.1,    # 标签平滑
         "mosaic": True,            # 启用Mosaic增强
-        # "class_weights": [1.0, 2.0]  # 设置类别权重，insulator权重更高
+        # "class_weights": [2.0, 2.0, 1.0]  # 设置类别权重
     },
     "stage2": {
         "epochs": 150,              # 增加训练轮数