code: def test_scheduler(device):
      import torch
      from pytest import raises
      from torch.optim import SGD
      from speechbrain.core import Brain
      from speechbrain.utils.profiling import profile, schedule
  
      @schedule
      @profile
      class SimpleBrain(Brain):
          def compute_forward(self, batch, stage):
              return self.modules.model(batch[0])
  
          def compute_objectives(self, predictions, batch, stage):
              return torch.nn.functional.l1_loss(predictions, batch[1])
  
      model = torch.nn.Linear(
          in_features=10, out_features=10, device=device
      )
      inputs = torch.rand(10, 10, device=device)
      targets = torch.rand(10, 10, device=device)
      train_set = ([inputs, targets],)
      valid_set = ([inputs, targets],)
      test_set = ([inputs, targets],)
  
      # Profiling: __init__ constructor -- while scheduler: waiting --> nothing to report
      brain = SimpleBrain(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      assert brain.profiler.profiler is None
      assert len(brain.profiler.speechbrain_event_traces) == 0
      with raises(Exception) as err:
          brain.profiler.events()  # Tracing hasn't started, yet, so everything is in err. Scheduler says: wait.
      assert err.type == AssertionError
      assert brain.profiler.step_num == 0
  
      # Profiling: fit() for train operations.
      brain.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
      assert brain.profiler.step_num == 20
      assert len(brain.profiler.speechbrain_event_traces) == 1
      # assert len(brain.profiler.events()) >= 250  # 293 with torch==1.10.1
      # assert len(brain.profiler.key_averages()) >= 60  # 73 with torch==1.10.1
      """print(brain.profiler.key_averages().table(sort_by="cpu_time_total"))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg    # of Calls
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                ProfilerStep*        55.48%       1.504ms        99.00%       2.684ms       1.342ms             2
                                                aten::l1_loss         1.07%      29.000us         9.30%     252.000us      63.000us             4
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         4.57%     124.000us         7.01%     190.000us      47.500us             4
                                                 aten::linear         2.32%      63.000us         6.93%     188.000us      94.000us             2
                                               aten::isfinite         0.89%      24.000us         5.35%     145.000us      48.333us             3
                                                     aten::to         1.00%      27.000us         4.46%     121.000us       7.562us            16
                                               aten::_to_copy         2.18%      59.000us         3.58%      97.000us       9.700us            10
                                                  aten::stack         0.89%      24.000us         3.28%      89.000us      17.800us             5
                                       aten::l1_loss_backward         0.55%      15.000us         3.14%      85.000us      42.500us             2
                                                   aten::mean         0.48%      13.000us         3.10%      84.000us      42.000us             2
                                                 aten::matmul         0.55%      15.000us         2.84%      77.000us      38.500us             2
                                                     aten::ne         0.77%      21.000us         2.73%      74.000us      24.667us             3
                                                     aten::mm         2.29%      62.000us         2.36%      64.000us      21.333us             3
                                                    aten::div         1.36%      37.000us         2.25%      61.000us      20.333us             3
                                      Optimizer.step#SGD.step         1.88%      51.000us         2.21%      60.000us      60.000us             1
      autograd::engine::evaluate_function: L1LossBackward0...         0.11%       3.000us         2.07%      56.000us      56.000us             1
                                                  aten::zeros         1.44%      39.000us         1.99%      54.000us       6.750us             8
                                              L1LossBackward0         0.11%       3.000us         1.95%      53.000us      53.000us             1
                                                      aten::t         0.89%      24.000us         1.88%      51.000us      10.200us             5
                                                    aten::cat         0.59%      16.000us         1.81%      49.000us       9.800us             5
                                                   aten::div_         0.70%      19.000us         1.73%      47.000us      23.500us             2
             autograd::engine::evaluate_function: MmBackward0         0.11%       3.000us         1.70%      46.000us      46.000us             1
                                                    aten::abs         1.14%      31.000us         1.59%      43.000us       5.375us             8
                                                  MmBackward0         0.37%      10.000us         1.59%      43.000us      43.000us             1
                                                    aten::sum         1.00%      27.000us         1.33%      36.000us      12.000us             3
                                                  aten::empty         1.29%      35.000us         1.29%      35.000us       1.207us            29
            autograd::engine::evaluate_function: AddBackward0         0.63%      17.000us         1.25%      34.000us      34.000us             1
                                                   aten::_cat         0.85%      23.000us         1.22%      33.000us       6.600us             5
                                                   aten::add_         1.11%      30.000us         1.11%      30.000us       5.000us             6
                                              aten::transpose         0.70%      19.000us         0.92%      25.000us       5.000us             5
            autograd::engine::evaluate_function: DivBackward0         0.18%       5.000us         0.92%      25.000us      25.000us             1
                                                   aten::norm         0.77%      21.000us         0.89%      24.000us       8.000us             3
                            Optimizer.zero_grad#SGD.zero_grad         0.63%      17.000us         0.81%      22.000us      22.000us             1
                                                   aten::item         0.55%      15.000us         0.77%      21.000us       3.000us             7
                                                  aten::copy_         0.77%      21.000us         0.77%      21.000us       2.100us            10
                                                 DivBackward0         0.15%       4.000us         0.74%      20.000us      20.000us             1
      autograd::engine::evaluate_function: torch::autograd...         0.15%       4.000us         0.74%      20.000us      10.000us             2
                                                    aten::mul         0.55%      15.000us         0.74%      20.000us       5.000us             4
                                                 aten::detach         0.37%      10.000us         0.70%      19.000us       3.167us             6
                                             aten::as_strided         0.59%      16.000us         0.63%      17.000us       1.062us            16
                                              aten::unsqueeze         0.41%      11.000us         0.59%      16.000us       2.667us             6
                                          aten::empty_strided         0.59%      16.000us         0.59%      16.000us       1.455us            11
                              torch::autograd::AccumulateGrad         0.30%       8.000us         0.59%      16.000us       8.000us             2
                                             aten::is_nonzero         0.18%       5.000us         0.59%      16.000us       5.333us             3
                                                   aten::view         0.55%      15.000us         0.55%      15.000us       3.000us             5
                                                aten::random_         0.52%      14.000us         0.52%      14.000us       7.000us             2
      autograd::engine::evaluate_function: UnsafeViewBackw...         0.11%       3.000us         0.52%      14.000us      14.000us             1
                                                    aten::sub         0.48%      13.000us         0.48%      13.000us       4.333us             3
                                                    aten::add         0.18%       5.000us         0.48%      13.000us      13.000us             1
                                              aten::ones_like         0.15%       4.000us         0.44%      12.000us      12.000us             1
                                                   aten::mul_         0.44%      12.000us         0.44%      12.000us       4.000us             3
                                                       detach         0.33%       9.000us         0.44%      12.000us       2.000us             6
                                          UnsafeViewBackward0         0.11%       3.000us         0.41%      11.000us      11.000us             1
                                                   aten::abs_         0.18%       5.000us         0.37%      10.000us       5.000us             2
                                             aten::empty_like         0.26%       7.000us         0.37%      10.000us       5.000us             2
                                                  aten::clamp         0.26%       7.000us         0.37%      10.000us      10.000us             1
                                             aten::zeros_like         0.18%       5.000us         0.33%       9.000us       9.000us             1
                                                     aten::eq         0.33%       9.000us         0.33%       9.000us       3.000us             3
                                                  aten::zero_         0.30%       8.000us         0.30%       8.000us       0.727us            11
                                           aten::_unsafe_view         0.22%       6.000us         0.30%       8.000us       4.000us             2
                                                  aten::fill_         0.30%       8.000us         0.30%       8.000us       2.000us             4
                                                aten::reshape         0.18%       5.000us         0.30%       8.000us       8.000us             1
                                    aten::_local_scalar_dense         0.22%       6.000us         0.26%       7.000us       1.000us             7
              autograd::engine::evaluate_function: TBackward0         0.11%       3.000us         0.26%       7.000us       7.000us             1
                                                aten::resize_         0.22%       6.000us         0.22%       6.000us       1.200us             5
                                             aten::reciprocal         0.22%       6.000us         0.22%       6.000us       6.000us             1
                                                   aten::sgn_         0.15%       4.000us         0.15%       4.000us       4.000us             1
                                                   TBackward0         0.07%       2.000us         0.15%       4.000us       4.000us             1
                                         aten::_reshape_alias         0.11%       3.000us         0.11%       3.000us       3.000us             1
                                              aten::clamp_max         0.11%       3.000us         0.11%       3.000us       3.000us             1
                                           aten::resolve_conj         0.07%       2.000us         0.07%       2.000us       0.333us             6
                                      aten::broadcast_tensors         0.07%       2.000us         0.07%       2.000us       1.000us             2
                                                 AddBackward0         0.04%       1.000us         0.04%       1.000us       1.000us             1
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 2.711ms  <===  above: Self CPU time total: 18.451ms (... the impact of warm-up)
      """
  
      @schedule
      @profile
      def train():
          # The step() function is executed inside speechbrain.core.brain.fit and is property of the Brain's profiler.
          # Above profiler and its scheduler are without power here, since prof.step() is not run - at all.
          brain.fit(
              epoch_counter=range(10),
              train_set=train_set,
              valid_set=valid_set,
          )
  
      prof = train()
      # since we used the same brain (which has its own profiler)
      # assert brain.profiler.step_num == 20  # started again from 0 steps
      assert len(brain.profiler.speechbrain_event_traces) == 2
      # assert len(brain.profiler.events()) >= 250  # 293 with torch==1.10.1  # unchanged (overwritten with akin data)
      # assert len(brain.profiler.key_averages()) >= 60  # 73 with torch==1.10.1 # unchanged (akin data)
      # now, to the train function's profiler
      assert (
          prof.step_num == 0
      )  # the prof.step() operation wasn't run (not in scope) -> its scheduler is unawaken!
      assert not hasattr(
          prof, "speechbrain_event_traces"
      )  # no trace collection
      with raises(Exception) as err_prof:
          prof.events()  # No tracing started with this one.
      assert err_prof.type == AssertionError  # sparing: key_averages()
  
      # But how to add profiling then if no writing access is there for a class... pretrained, for example:
      class SimpleBrainUntracked(Brain):
          def compute_forward(self, batch, stage):
              return self.modules.model(batch[0])
  
          def compute_objectives(self, predictions, batch, stage):
              return torch.nn.functional.l1_loss(predictions, batch[1])
  
      brain_or_pretrained = SimpleBrainUntracked(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
  
      # Set-up the profiler and hook it to the model.
      scheduled_profiler = schedule(profile)
      scheduled_profiler(brain_or_pretrained)
  
      # Profiling: still too early for scheduler!
      brain_or_pretrained.evaluate(test_set=test_set)  # -> step_num=1
      assert brain_or_pretrained.profiler.step_num == 1
      assert brain_or_pretrained.profiler.profiler is None
  
      # Profiling: scheduler warms-up.
      brain_or_pretrained.evaluate(
          test_set=(
              [inputs, targets],  # +1x test_set -> step_num=2
              [inputs, targets],
          )
      )
      assert brain_or_pretrained.profiler.step_num == 2
      # brain_or_pretrained.profiler.profiler will be set (not None anymore)
      # when run on cpu, there are no events - but cuda activities are recorded if existing
      # see: https://github.com/speechbrain/speechbrain/issues/1469
      if (
          torch.profiler.ProfilerActivity.CUDA
          in brain_or_pretrained.profiler.activities
      ):
          assert (
              len(
                  set(
                      [
                          x.name
                          for x in brain_or_pretrained.profiler.profiler.function_events
                      ]
                  )
                  - {
                      "cudaGetDeviceCount",
                      "cudaGetDeviceProperties",
                      "cudaDeviceSynchronize",
                  }
              )
          ) == 0
      else:
          assert len(brain_or_pretrained.profiler.events()) == 0
  
      # Profiling: scheduler warms-up...
      brain_or_pretrained.evaluate(
          test_set=(
              [inputs, targets],  # +1x test_set
              [inputs, targets],  # +2x test_set -> step_num=3
              [inputs, targets],
          )
      )
      assert brain_or_pretrained.profiler.step_num == 3
      if (
          torch.profiler.ProfilerActivity.CUDA
          in brain_or_pretrained.profiler.activities
      ):
          assert (
              len(
                  set(
                      [
                          x.name
                          for x in brain_or_pretrained.profiler.profiler.function_events
                      ]
                  )
                  - {
                      "cudaGetDeviceCount",
                      "cudaGetDeviceProperties",
                      "cudaDeviceSynchronize",
                  }
              )
          ) == 0
      else:
          assert len(brain_or_pretrained.profiler.events()) == 0
  
      # Profiling: first trace!
      brain_or_pretrained.evaluate(
          test_set=(
              [inputs, targets],  # +1x test_set
              [inputs, targets],  # +2x test_set
              [inputs, targets],  # +3x test_set -> step_num=4
              [inputs, targets],
          )
      )
      assert brain_or_pretrained.profiler.step_num == 4
      # assert len(brain_or_pretrained.profiler.events()) >= 4  # == 10  # before
      # assert len(brain_or_pretrained.profiler.key_averages()) >= 4  # == 5  # before
      assert len(brain_or_pretrained.profiler.events()) >= 1
  
label_desc: Use image node_id 3 for calling a global function with example usage: test_scheduler(device) and returns: self, torch, self, torch