code: def prepare_profiler_for_brain(prof: profiler.profile):
      """Sets up a ``torch.profiler.profile`` to also (a) aggregate traces issued from various interactions
      with ``speechbrain.core.Brain``:s and (b) hooks a method to ``merge_traces``.
      """
      # Brain functions will be called independently -> traces will be segregated, so we aggregate them.
      prof.speechbrain_event_traces = list()
  
      # Preparing the profiler to be re-used during Brain:s' lifecycles.
      def hook_profiler_stop(stop: Callable):
          """Implementation of hook_profiler_stop."""
  
          @wraps(stop)
          def stop_wrapper():
              """Implementation of stop_wrapper."""
              kineto_profiler = prof.profiler
              if kineto_profiler is not None:
                  stop_result = stop()
                  if (
                      prof.events()
                  ):  # kineto events are not aggregatable (sticking with parsed kineto events)
                      # see: torch.autograd.profiler.__exit__
                      kineto_events = (
                          kineto_profiler._parse_kineto_results(
                              kineto_profiler.kineto_results
                          )
                      )
                      # add to trace record
                      prof.speechbrain_event_traces.append(
                          deepcopy(kineto_events)
                      )
                      # set flag to disable the profiler
                      kineto_profiler.enabled = False
                  return stop_result
              else:
                  return stop()  # will be: None
  
          return stop_wrapper
  
      # Preparing the profiler to be re-started during Brain:s' lifecycles.
      def hook_profiler_start(start: Callable):
          """Implementation of hook_profiler_start."""
  
          @wraps(start)
          def start_wrapper():
              """Implementation of start_wrapper."""
              prof.step_num = 0
              prof.current_action = prof.schedule(prof.step_num)
              kineto_profiler = prof.profiler
              if kineto_profiler is not None:
                  # check flag if profiler is disabled (i.e. as of stop_wrapper); prevents entering its __init__ twice
                  if not kineto_profiler.enabled:
                      # reset kineto profiler (otherwise, one obtains the same traces over & over again)
                      kineto_profiler.enabled = True
              return start()
  
          return start_wrapper
  
      """
      # It's currently designed as hiding an Easter Egg.
      def merge_traces():
          " ""Implementation of merge_traces." ""
          # Alternative re-design quirks: make trace aggregator a GLOBAL -or- create another profiler class.
          trace_aggregator = "speechbrain_event_traces"
          if prof.profiler is not None:
              if trace_aggregator in dir(prof) and prof.events():
                  # clear all assigned parents/children (from previous mergers & trees)
                  for trace in getattr(prof, trace_aggregator):
                      for event in trace:
                          event.cpu_parent = None
                          event.cpu_children: List[FunctionEvent] = []
                  # assemble new list
                  merged_events = EventList(
                      list(chain.from_iterable(getattr(prof, trace_aggregator))),
                      use_cuda=prof.profiler.use_cuda,
                      profile_memory=prof.profiler.profile_memory,
                      with_flops=prof.profiler.with_flops,
                  )
                  merged_events._build_tree()
                  return merged_events
              else:  # not tested
                  return prof.events()
          else:
              return []
      """
  
      # Augment torch's profiler.
      setattr(
          prof, "start", hook_profiler_start(getattr(prof, "start"))
      )
      setattr(prof, "stop", hook_profiler_stop(getattr(prof, "stop")))
      # setattr(prof, "merge_traces", merge_traces)
  
      # Return so it can be readily assigned elsewhere :)
      return prof
  
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 4 for calling a global function with example usage: prepare_profiler_for_brain(prof) and returns: prof, stop_wrapper, start_wrapper, start, stop_result, stop