diff --git a/config/lr_scheduler.yaml b/config/lr_scheduler.yaml
index e021f03..2fdd612 100644
--- a/config/lr_scheduler.yaml
+++ b/config/lr_scheduler.yaml
@@ -3,7 +3,7 @@ scheduler_type: LINEAR_WARMUP_THEN_POLY_SCHEDULER
 # total_iters=epochs * 训练的图像数量
 kwargs: |
   {
-    "warmup_iters": 4356, 
+    "warmup_iters": 2904, 
     "total_iters": 290400, 
     "warmup_ratio": 0.000001, 
     "min_lr": 0., 
diff --git a/config/train.yaml b/config/train.yaml
index 5733fcc..945bdc9 100644
--- a/config/train.yaml
+++ b/config/train.yaml
@@ -8,6 +8,8 @@ epochs: 100
 # 每一类的占比权重，如果要让每一类的占比权重相同，为1.0即可
 weight:
   - 1.0
+  - 1.0
+  - 1.0
 
 # 数据集存放位置
 root:
diff --git a/losses.py b/losses.py
index 717b3ff..c9d5ca6 100644
--- a/losses.py
+++ b/losses.py
@@ -13,7 +13,8 @@ class FocalLoss(nn.Module):
         super(FocalLoss, self).__init__()
         self.gamma = gamma
         self.eps = eps
-        self.ce = torch.nn.CrossEntropyLoss(weight=weight, reduction=reduction)
+        myweight = torch.tensor([1.0, 1.9, 1.0, 1.2]).cuda()
+        self.ce = torch.nn.CrossEntropyLoss(weight=myweight, reduction=reduction)
 
     def forward(self, x, y):
         logp = self.ce(x, y)
diff --git a/test.ipynb b/test.ipynb
index fcd37e5..5dc1bf9 100644
--- a/test.ipynb
+++ b/test.ipynb
@@ -2,17 +2,18 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
     "from PIL import Image\n",
-    "import numpy as np"
+    "import numpy as np\n",
+    "import os"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [
     {
@@ -84,6 +85,13 @@
     "print(image_array.shape)\n",
     "print(image_array)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {