code: def profile(
      func: Optional[object] = None,
      class_hooks: Optional[Iterable[str]] = None,
      activities: Optional[Iterable[profiler.ProfilerActivity]] = None,
      schedule: Optional[
          Callable[[int], profiler.ProfilerAction]
      ] = None,
      on_trace_ready: Optional[Callable[..., Any]] = None,
      record_shapes: bool = False,
      profile_memory: bool = False,
      with_stack: bool = False,
      with_flops: bool = False,
      with_modules: bool = False,
  ) -> object:
      """Wrapper to create a PyTorch profiler to benchmark training/inference of speechbrain.core.Brain instances.
      See ``torch.profiler.profile`` documentation for details (brief summary below).
  
      Arguments
      ---------
      func : object
          ``speechbrain.core.Brain``:s or a (train/eval) function to be profiled.
      class_hooks : iterable
          List of method/function names of ``speechbrain.core.Brain``:s that should be profiled also.
          Otherwise, only the __init__ constructor will be profiled when decorating a Brain class.
          Default: ``['fit', 'evaluate']`` for classes, and ``None`` for functions.
      activities : iterable
          List of activity groups.
          Default: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA.
          (Default value should be ok for most cases.)
      schedule : callable
          Waits a specified amount of steps for PyTorch to warm-up; see the above ``schedule`` decorator.
          Default: ``ProfilerAction.RECORD`` (immediately starts recording).
      on_trace_ready : callable
          Specifies what benchmark record should be saved (after each scheduled step);
          see above ``trace_handler`` decorator.
          Default: ``None`` (pick up collected reporting once profiling ended, but not details per step).
      record_shapes : bool
          Save input shapes of operations (enables to group benchmark data by after profiling).
          Default: ``False``.
      profile_memory : bool
          Track tensor memory allocation/deallocation.
          Default: ``False``.
      with_stack : bool
          Record source information (file and line number).
          Default: ``False``.
      with_flops: bool
          Estimate the number of FLOPs.
          Default: ``False``.
      with_modules: bool
          Record module hierarchy (including function names)
          Default: ``False``
  
      Example
      -------
      >>> import torch
      >>> @profile
      ... def run(x : torch.Tensor):
      ...     y = x ** 2
      ...     z = y ** 3
      ...     return y.backward()  # y.backward() returns None --> return value is substituted with profiler
      >>> data = torch.randn((1, 1), requires_grad=True)
      >>> prof = run(data)
      >>> out = [len(prof.events()), len(prof.key_averages()), prof.profiler.total_average().count]
      """
      if func is None:  # return a profiler; not tested
          return prepare_profiler_for_brain(
              profiler.profile(
                  activities=activities,
                  schedule=schedule,
                  on_trace_ready=on_trace_ready,
                  record_shapes=record_shapes,
                  profile_memory=profile_memory,
                  with_stack=with_stack,
                  with_flops=with_flops,
                  with_modules=with_modules,
              )
          )
      # Polymorph: func is pretrained or an instance of Brain (assumed case)
      if hasattr(func, "HPARAMS_NEEDED") or not callable(func):
          with profiler.profile(
              activities=activities,
              schedule=schedule,  # scheduler needs to be set directly (fetching is here not possible as for wrappers)
              on_trace_ready=on_trace_ready,
              record_shapes=record_shapes,
              profile_memory=profile_memory,
              with_stack=with_stack,
              with_flops=with_flops,
              with_modules=with_modules,
          ) as prof:
              func.profiler = prepare_profiler_for_brain(prof)
              hook_brain_methods(
                  func=func, class_hooks=class_hooks, prof=prof
              )
              return func  # no need to return anything; all done in-place; but if needs to be readily assigned elsewhere
      else:
          # callable(func) - polymorph: __init__ Brain constructor -or- function to be wrapped
          @wraps(func)
          def wrapper(*args, **kwargs):
              """Implementation of the wrapper."""
              # Binding variables.
              nonlocal class_hooks
              nonlocal schedule
              nonlocal on_trace_ready
              # Check if there's a nested decorators.
              if schedule is None:
                  if "schedule" in kwargs:
                      schedule = kwargs.pop("schedule")
              if on_trace_ready is None:
                  if "on_trace_ready" in kwargs:
                      on_trace_ready = kwargs.pop("on_trace_ready")
              with profiler.profile(
                  activities=activities,
                  schedule=schedule,
                  on_trace_ready=on_trace_ready,
                  record_shapes=record_shapes,
                  profile_memory=profile_memory,
                  with_stack=with_stack,
                  with_flops=with_flops,
                  with_modules=with_modules,
              ) as prof:
                  # Preserves profiler as class attribute if func is not a function (implies: speechbrain.core.Brain).
                  if "__call__" not in dir(func):
                      # Passing the profiler to Bain:s' __init__ constructor as an additional argument.
                      kwargs["profiler"] = prepare_profiler_for_brain(
                          prof
                      )
                      hook_brain_methods(
                          func=func, class_hooks=class_hooks, prof=prof
                      )
  
                  # Run & trace to benchmark.
                  result = func(*args, **kwargs)
  
                  # Prof is about to be lost at return.
                  if "__call__" in dir(func):
                      if result is None:
                          return prof  # for void function, simply return profiling data
                      else:  # not tested - returns both
                          return result, prof
  
                  return result
  
          return wrapper
  
file_path: speechbrain/speechbrain/utils/profiling.py
imports: import numpy
  from copy import deepcopy
  from torch import profiler
  from functools import wraps
  from typing import Any, Callable, Iterable, Optional
label_desc: Use image node_id 6 for calling a global function with example usage: profile(func, class_hooks, activities, schedule, on_trace_ready, record_shapes, profile_memory, with_stack, with_flops, with_modules) and returns: prepare_profiler_for_brain, wrapper, func, result, prof, result,  prof