About mghead loss compute question about det3d HOT 5 CLOSED

if not test:
            annotations = [
                nusc.get("sample_annotation", token) for token in sample["anns"]
            ]

            locs = np.array([b.center for b in ref_boxes]).reshape(-1, 3)
            dims = np.array([b.wlh for b in ref_boxes]).reshape(-1, 3)
            # rots = np.array([b.orientation.yaw_pitch_roll[0] for b in ref_boxes]).reshape(-1, 1)
            # velocity = np.array([b.velocity for b in ref_boxes]).reshape(-1, 3)
            velocity = np.array(
                [nusc.box_velocity(token)[:2] for token in sample['anns']]
            )
            # convert velo from global to lidar
            for i in range(len(ref_boxes)):
                velo = np.array([*velocity[i], 0.0])
                velo = velo @ np.linalg.inv(e2g_r_mat).T @ np.linalg.inv(
                    l2e_r_mat).T
                velocity[i] = velo[:2]
            velocity = velocity.reshape(-1,2)

            rots = np.array([quaternion_yaw(b.orientation) for b in ref_boxes]).reshape(
                -1, 1
            )
            names = np.array([b.name for b in ref_boxes])
            tokens = np.array([b.token for b in ref_boxes])
            gt_boxes = np.concatenate(
                [locs, dims, velocity[:, :2], -rots - np.pi / 2], axis=1
            )

if kwargs.get("mode", False):
    reg_targets = example["reg_targets"][task_id][:, :, [0, 1, 3, 4, 6]]
    reg_targets_left = example["reg_targets"][task_id][:, :, [2, 5]]
else:
    reg_targets = example["reg_targets"][task_id]

## Add part
for i in range(6):
    example["reg_targets"][i][torch.isnan(example["reg_targets"][i])] = 0.0