class_docstr: Summary class for a single layer in a :class:`~lightning.pytorch.core.LightningModule`. It collects the
  following information:
  
  - Type of the layer (e.g. Linear, BatchNorm1d, ...)
  - Input shape
  - Output shape
  - Number of parameters
  
  The input and output shapes are only known after the example input array was
  passed through the model.
  
  Example::
  
      >>> model = torch.nn.Conv2d(3, 8, 3)
      >>> summary = LayerSummary(model)
      >>> summary.num_parameters
      224
      >>> summary.layer_type
      'Conv2d'
      >>> output = model(torch.rand(1, 3, 5, 5))
      >>> summary.in_size
      [1, 3, 5, 5]
      >>> summary.out_size
      [1, 8, 3, 3]
  
  Args:
      module: A module to summarize
class_docstr_tok: ["Summary","class","for","a","single","layer","in","a",":","class",":","`","~lightning.pytorch.core.LightningModule","`",".","It","collects","the","following","information",":","-","Type","of","the","layer","(","e.g",".","Linear",",","BatchNorm1d",",","...",")","-","Input","shape","-","Output","shape","-","Number","of","parameters","The","input","and","output","shapes","are","only","known","after","the","example","input","array","was","passed","through","the","model",".","Example",":",":",">",">",">","model","=","torch.nn.Conv2d","(","3",",","8",",","3",")",">",">",">","summary","=","LayerSummary","(","model",")",">",">",">","summary.num_parameters","224",">",">",">","summary.layer_type","'Conv2d'",">",">",">","output","=","model","(","torch.rand","(","1",",","3",",","5",",","5",")",")",">",">",">","summary.in_size","[","1",",","3",",","5",",","5","]",">",">",">","summary.out_size","[","1",",","8",",","3",",","3","]","Args",":","module",":","A","module","to","summarize"]
docstr: 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.
docstr_tok: ["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","."]
code: 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):
          if _TORCH_GREATER_EQUAL_2_0:
              handle = self._module.register_forward_hook(
                  hook_with_kwargs, with_kwargs=True
              )
          else:
              handle = self._module.register_forward_hook(hook)
  
      return handle
  
imports: import contextlib
  import logging
  import math
  from collections import OrderedDict
  from typing import Any, Dict, List, Optional, Tuple, Union
  import torch
  import torch.nn
  from torch import Tensor
  from torch.utils.hooks import RemovableHandle
  import lightning.pytorch
  from lightning.fabric.utilities.imports import _TORCH_GREATER_EQUAL_2_0
  from lightning.pytorch.utilities.model_helpers import _ModuleMode
  from lightning.pytorch.utilities.rank_zero import WarningCache
label_desc: Use image node_id 3 for calling the LayerSummary obj's underlying member method code with example usage: obj._register_hook() and returns: handle