fix planning cumsum dimension

OpenDriveLab · Aug 7, 2023 · 7903ead · 7903ead
1 parent 489a297
commit 7903ead
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/projects/mmdet3d_plugin/uniad/dense_heads/planning_head.py b/projects/mmdet3d_plugin/uniad/dense_heads/planning_head.py
@@ -188,7 +188,7 @@ def forward(self,
         plan_query = self.attn_module(plan_query, bev_feat)   # [1, 1, 256]
 
         sdc_traj_all = self.reg_branch(plan_query).view((-1, self.planning_steps, 2))
-        sdc_traj_all[...,:2] = torch.cumsum(sdc_traj_all[...,:2], dim=2)
+        sdc_traj_all[...,:2] = torch.cumsum(sdc_traj_all[...,:2], dim=1)
         sdc_traj_all[0] = bivariate_gaussian_activation(sdc_traj_all[0])
         if self.use_col_optim and not self.training:
             # post process, only used when testing