Repository: Project-MONAI/MONAI
File: monai/apps/detection/networks/retinanet_network.py (Python)

Code before:
        self.spatial_dims = look_up_option(spatial_dims, supported=[1, 2, 3])
        self.num_classes = num_classes
        self.size_divisible = ensure_tuple_rep(size_divisible, self.spatial_dims)
        self.use_list_output = use_list_output

        if not hasattr(feature_extractor, "out_channels"):
            raise ValueError(
                "feature_extractor should contain an attribute out_channels "
                "specifying the number of output channels (assumed to be the "
                "same for all the levels)"
            )
        self.feature_extractor = feature_extractor

        self.feature_map_channels: int = self.feature_extractor.out_channels  # type: ignore[assignment]
        self.num_anchors = num_anchors
        self.classification_head = RetinaNetClassificationHead(
            self.feature_map_channels, self.num_anchors, self.num_classes, spatial_dims=self.spatial_dims
        )
        self.regression_head = RetinaNetRegressionHead(
            self.feature_map_channels, self.num_anchors, spatial_dims=self.spatial_dims
        )

        self.cls_key: str = "classification"
        self.box_reg_key: str = "box_regression"

    def forward(self, images: Tensor) -> dict[str, list[Tensor]] | list[Tensor]:
        """
        It takes an image tensor as inputs, and outputs predicted classification maps
        and predicted box regression maps in ``head_outputs``.

        Args:
            images: input images, sized (B, img_channels, H, W) or (B, img_channels, H, W, D).

        Return:
            1) If self.use_list_output is False, output a dictionary ``head_outputs`` with
            keys including self.cls_key and self.box_reg_key.
            ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
            ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
            2) if self.use_list_output is True, outputs a list of 2N tensors ``head_outputs``, with first N tensors being
            the predicted classification maps and second N tensors being the predicted box regression maps.

        """
        # compute features maps list from the input images.
        features = self.feature_extractor(images)
        if isinstance(features, Tensor):
            feature_maps = [features]
        elif torch.jit.isinstance(features, Dict[str, Tensor]):
            feature_maps = list(features.values())
        else:
            feature_maps = list(features)


Review comment (suggestion):
Hi @Can-Zhao ,
    The TorchScript doesn't support the type hint with `|` in the `forward` method (other methods are fine).  It will report  an error below. Could you please replace it with a `Union` or `Any`. `Union` would be better if it can pass the ci test.

```
Expression of type | cannot be used in a type expression:
  File "/home/liubin/data/github_monai/det-can/MONAI/monai/apps/detection/networks/retinanet_network.py", line 305
    def forward(self, images: Tensor) -> dict[str, list[Tensor]] | list[Tensor]:
                                         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
        """
        It takes an image tensor as inputs, and outputs predicted classification maps
```

Thanks,
Bin

Code after:
        self.spatial_dims = look_up_option(spatial_dims, supported=[1, 2, 3])
        self.num_classes = num_classes
        self.size_divisible = ensure_tuple_rep(size_divisible, self.spatial_dims)
        self.use_list_output = use_list_output

        if not hasattr(feature_extractor, "out_channels"):
            raise ValueError(
                "feature_extractor should contain an attribute out_channels "
                "specifying the number of output channels (assumed to be the "
                "same for all the levels)"
            )
        self.feature_extractor = feature_extractor

        self.feature_map_channels: int = self.feature_extractor.out_channels  # type: ignore[assignment]
        self.num_anchors = num_anchors
        self.classification_head = RetinaNetClassificationHead(
            self.feature_map_channels, self.num_anchors, self.num_classes, spatial_dims=self.spatial_dims
        )
        self.regression_head = RetinaNetRegressionHead(
            self.feature_map_channels, self.num_anchors, spatial_dims=self.spatial_dims
        )

        self.cls_key: str = "classification"
        self.box_reg_key: str = "box_regression"

    def forward(self, images: Tensor) -> Any:
        """
        It takes an image tensor as inputs, and outputs predicted classification maps
        and predicted box regression maps in ``head_outputs``.

        Args:
            images: input images, sized (B, img_channels, H, W) or (B, img_channels, H, W, D).

        Return:
            1) If self.use_list_output is False, output a dictionary ``head_outputs`` with
            keys including self.cls_key and self.box_reg_key.
            ``head_outputs[self.cls_key]`` is the predicted classification maps, a list of Tensor.
            ``head_outputs[self.box_reg_key]`` is the predicted box regression maps, a list of Tensor.
            2) if self.use_list_output is True, outputs a list of 2N tensors ``head_outputs``, with first N tensors being
            the predicted classification maps and second N tensors being the predicted box regression maps.

        """
        # compute features maps list from the input images.
        features = self.feature_extractor(images)
        if isinstance(features, Tensor):
            feature_maps = [features]
        elif torch.jit.isinstance(features, Dict[str, Tensor]):
            feature_maps = list(features.values())
        else:
            feature_maps = list(features)