Repository: Project-MONAI/MONAI
File: monai/inferers/inferer.py (Python)

Code before:
        if isinstance(out, torch.Tensor):
            return out.unsqueeze(dim=self.spatial_dim + 2)

        if isinstance(out, Mapping):
            for k in out.keys():
                out[k] = out[k].unsqueeze(dim=self.spatial_dim + 2)
            return out

        return tuple(out_i.unsqueeze(dim=self.spatial_dim + 2) for out_i in out)


class SlidingWindowInfererAdapt(SlidingWindowInferer):
    """
    SlidingWindowInfererAdapt extends SlidingWindowInferer to automatically switch to CPU stitching,
    when GPU stitching runs out of GPU memory. It also records a size of such large images to automatically
    try CPU stitching for the next large image of a similar size.  If the stitching 'device' input parameter is provided,
    automatic switch to CPU won't be attempted.

    """

    def __call__(
        self,
        inputs: torch.Tensor,
        network: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]],
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]:
        """

        Args:
            inputs: model input data for inference.
            network: target model to execute inference.
                supports callables such as ``lambda x: my_torch_model(x, additional_config)``
            args: optional args to be passed to ``network``.
            kwargs: optional keyword args to be passed to ``network``.

        """

        try:
            out = super().__call__(inputs=inputs, network=network, *args, **kwargs)
        except RuntimeError as e:

            if (
                inputs.is_cuda
                and self.device is None
                and not (self.cpu_thresh is not None and inputs.shape[2:].numel() > self.cpu_thresh)
                and "OutOfMemoryError" in str(type(e).__name__)
            ):

                print(e)
                warnings.warn(f"GPU stitching failed, attempting CPU stitching, image dim {inputs.shape}")

Review comment (bug):
this could be useful for all the `Inferer` subclasses, how about:
if you agree I can update the PR to include this

```py
class AutoDeviceInferer(Inferer):
    def __init__(self, inferer):
        self.inferer = inferer
        self.cpu_thresh = getattr(self.inferer, "cpu_thresh", None)

    def __call__(
        self,
        inputs: torch.Tensor,
        network: Callable[..., torch.Tensor | Sequence[torch.Tensor] | dict[Any, torch.Tensor]],
        *args: Any,
        **kwargs: Any,
    ) -> torch.Tensor | tuple[torch.Tensor, ...] | dict[Any, torch.Tensor]:
        try:
            return self.inferer(inputs=inputs, network=network, *args, **kwargs)
        except RuntimeError as e:
            if (
                inputs.is_cuda
                and torch.device(self.inferer.device).type == "cuda"
                and not (self.cpu_thresh is not None and inputs.shape[2:].numel() > self.cpu_thresh)
                and "OutOfMemoryError" in str(type(e).__name__)
            ):
                warnings.warn(f"GPU stitching failed {e}, attempting CPU stitching, image dim {inputs.shape}")
                self.cpu_thresh = inputs.shape[2:].numel() - 1  # remember large image size
                # CPU stitching (trigger cpu_thresh)
                return self.inferer(inputs=inputs, network=network, *args, **kwargs)
            raise e
```

Code after:
        kwargs: other optional keyword args to be passed to `__init__` of cam.

    """

    def __init__(
        self, cam_name: str, target_layers: str, class_idx: int | None = None, *args: Any, **kwargs: Any
    ) -> None:
        Inferer.__init__(self)
        if cam_name.lower() not in ("cam", "gradcam", "gradcampp"):
            raise ValueError("cam_name should be: 'CAM', 'GradCAM' or 'GradCAMpp'.")
        self.cam_name = cam_name.lower()
        self.target_layers = target_layers
        self.class_idx = class_idx
        self.args = args
        self.kwargs = kwargs

    def __call__(self, inputs: torch.Tensor, network: nn.Module, *args: Any, **kwargs: Any):  # type: ignore
        """Unified callable function API of Inferers.

        Args:
            inputs: model input data for inference.
            network: target model to execute inference.
                supports callables such as ``lambda x: my_torch_model(x, additional_config)``
            args: other optional args to be passed to the `__call__` of cam.
            kwargs: other optional keyword args to be passed to `__call__` of cam.

        """
        cam: CAM | GradCAM | GradCAMpp
        if self.cam_name == "cam":
            cam = CAM(network, self.target_layers, *self.args, **self.kwargs)
        elif self.cam_name == "gradcam":
            cam = GradCAM(network, self.target_layers, *self.args, **self.kwargs)
        else:
            cam = GradCAMpp(network, self.target_layers, *self.args, **self.kwargs)

        return cam(inputs, self.class_idx, *args, **kwargs)


class SliceInferer(SlidingWindowInferer):
    """
    SliceInferer extends SlidingWindowInferer to provide slice-by-slice (2D) inference when provided a 3D volume.
    A typical use case could be a 2D model (like 2D segmentation UNet) operates on the slices from a 3D volume,
    and the output is a 3D volume with 2D slices aggregated. Example::

        # sliding over the `spatial_dim`
        inferer = SliceInferer(roi_size=(64, 256), sw_batch_size=1, spatial_dim=1)
        output = inferer(input_volume, net)

    Args:
        spatial_dim: Spatial dimension over which the slice-by-slice inference runs on the 3D volume.
            For example ``0`` could slide over axial slices. ``1`` over coronal slices and ``2`` over sagittal slices.