[Docs] Incorrect type hint in the loss() function in mvx_two_stage.py

[ShaneXiangH]

Branch

main branch https://mmdetection3d.readthedocs.io/en/latest/

📚 The doc issue

This function:

def loss(self, batch_inputs_dict: Dict[List, torch.Tensor],
             batch_data_samples: List[Det3DDataSample],
             **kwargs) -> List[Det3DDataSample]:
        """
        Args:
            batch_inputs_dict (dict): The model input dict which include
                'points' and `imgs` keys.

                - points (list[torch.Tensor]): Point cloud of each sample.
                - imgs (torch.Tensor): Tensor of batch images, has shape
                  (B, C, H ,W)
            batch_data_samples (List[:obj:`Det3DDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance_3d`, .

        Returns:
            dict[str, Tensor]: A dictionary of loss components.

        """

        batch_input_metas = [item.metainfo for item in batch_data_samples]
        img_feats, pts_feats = self.extract_feat(batch_inputs_dict,
                                                 batch_input_metas)
        losses = dict()
        if pts_feats:
            losses_pts = self.pts_bbox_head.loss(pts_feats, batch_data_samples,
                                                 **kwargs)
            losses.update(losses_pts)
        if img_feats:
            losses_img = self.loss_imgs(img_feats, batch_data_samples)
            losses.update(losses_img)
        return losses

Shows incorrect type hint for the batch_inputs_dict param. It should be: Dict[str, Union[List[torch.Tensor], torch.Tensor]]

Suggest a potential alternative/fix

Change the type hint for the batch_inputs_dict param to: Dict[str, Union[List[torch.Tensor], torch.Tensor]]

[Xiangxu-0103]

Are you willing to create a PR to help us fix it?