diff --git a/README.md b/README.md
index bfa5ea9..f851797 100644
--- a/README.md
+++ b/README.md
@@ -154,7 +154,7 @@ cd YOPO/
 conda activate yopo
 python train_yopo.py
 ```
-It takes less than 1 hour to train on 100,000 samples for 50 epochs on an RTX 3080 GPU. Besides, we highly recommend binding the process to P-cores  via `taskset -c 1,2,3,4 python train_yopo.py` if your CPU uses a hybrid architecture with P-cores and E-cores. If everything goes well, the training log is as follows:
+It takes less than 1 hour to train on 100,000 samples for 50 epochs on an RTX 3080 GPU and i9-12900K CPU. Besides, we highly recommend binding the process to P-cores  via `taskset -c 1,2,3,4 python train_yopo.py` if your CPU uses a hybrid architecture with P-cores and E-cores. If everything goes well, the training log is as follows:
 
 ```
 cd YOPO/saved
diff --git a/YOPO/config/traj_opt.yaml b/YOPO/config/traj_opt.yaml
index 4994aae..7be5200 100644
--- a/YOPO/config/traj_opt.yaml
+++ b/YOPO/config/traj_opt.yaml
@@ -7,7 +7,7 @@ vel_max_train: 6.0
 acc_max_train: 6.0
 
 # IMPORTANT: weight of costs for unit speed (can be visualized in tensorboard)
-wg: 0.1    # guidance
+wg: 0.12   # guidance
 ws: 10.0   # smoothness
 wc: 0.1    # collision
 
diff --git a/YOPO/loss/safety_loss.py b/YOPO/loss/safety_loss.py
index e1310bc..cbe5f4f 100644
--- a/YOPO/loss/safety_loss.py
+++ b/YOPO/loss/safety_loss.py
@@ -57,10 +57,10 @@ class SafetyLoss(nn.Module):
         # get info from sdf_map
         cost = self.get_distance_cost(pos_batch, map_id)
 
-        cost_dt = (cost * dt).reshape(-1, pos_coe.shape[1])   # [B*H*V, N, 3]
-        cost_colli = cost_dt.sum(dim=-1, keepdim=True)
+        cost_dt = (cost * dt).reshape(-1, pos_coe.shape[1])   # [B*H*V, N]
+        cost_colli = cost_dt.sum(dim=-1)
 
-        return cost_colli.squeeze()
+        return cost_colli
 
     def get_distance_cost(self, pos, map_id):
         """