def _register_hook(self) -> Optional[RemovableHandle]:
        """Registers a hook on the module that computes the input- and output size(s) on the first forward pass. If the
        hook is called, it will remove itself from the from the module, meaning that recursive models will only record
        their input- and output shapes once. Registering hooks on :class:`~torch.jit.ScriptModule` is not supported.

        Return:
            A handle for the installed hook, or ``None`` if registering the hook is not possible.

        """

        def hook(_: nn.Module, inp: Any, out: Any) -> None:
            if len(inp) == 1:
                inp = inp[0]

            self._in_size = parse_batch_shape(inp)
            self._out_size = parse_batch_shape(out)
            assert self._hook_handle is not None
            self._hook_handle.remove()

        def hook_with_kwargs(_: nn.Module, args: Any, kwargs: Any, out: Any) -> None:
            # We can't write them in the same function, since the forward hook
            # uses positional arguments.

            inp = (*args, *kwargs.values()) if kwargs is not None else args
            hook(_, inp, out)

        handle = None
        if not isinstance(self._module, torch.jit.ScriptModule):
            handle = self._module.register_forward_hook(hook_with_kwargs, with_kwargs=True)

        return handle
Registers a hook on the module that computes the input- and output size(s) on the first forward pass. If the
        hook is called, it will remove itself from the from the module, meaning that recursive models will only record
        their input- and output shapes once. Registering hooks on :class:`~torch.jit.ScriptModule` is not supported.

        Return:
            A handle for the installed hook, or ``None`` if registering the hook is not possible.

        
src/lightning/pytorch/utilities/model_summary/model_summary.py