code: def test_profile_details(device):
      import torch
  
      # from copy import deepcopy
      from torch.optim import SGD
      from speechbrain.core import Brain
      from speechbrain.utils.profiling import (
          profile_analyst,
          profile_optimiser,
          export,
          # events_diff,
      )
  
      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],
          [inputs, targets],
          [inputs, targets],
          [inputs, targets],
          [inputs, targets],
          [inputs, targets],
      )
  
      brain_analyst = profile_analyst(
          SimpleBrain(
              {"model": model},
              lambda x: SGD(x, 0.1),
              run_opts={"device": device},
          )
      )
      brain_optimiser = profile_optimiser(
          SimpleBrain(
              {"model": model},
              lambda x: SGD(x, 0.1),
              run_opts={"device": device},
          )
      )
  
      assert len(brain_analyst.profiler.speechbrain_event_traces) == 0
      brain_analyst.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
      assert len(brain_analyst.profiler.speechbrain_event_traces) == 1
      # assert len(brain_analyst.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
      # assert len(brain_analyst.profiler.events()) >= 250  # 293 with torch==1.10.1
      # assert len(brain_analyst.profiler.events().key_averages()) >= 60  # 73 with torch==1.10.1
      """print(brain_analyst.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls   Total FLOPs
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                ProfilerStep*        36.62%       4.345ms        98.50%      11.686ms       5.843ms         796 b      -1.66 Kb             2            --
                                                aten::l1_loss         2.50%     297.000us        14.37%       1.705ms     426.250us          16 b        -800 b             4            --
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         3.91%     464.000us        11.01%       1.306ms     326.500us       1.55 Kb         -80 b             4            --
                                               aten::isfinite         3.10%     368.000us         9.65%       1.145ms     381.667us           3 b         -18 b             3            --
                                                  aten::stack         2.79%     331.000us         8.46%       1.004ms     200.800us       1.57 Kb           0 b             5            --
                                                     aten::to         1.96%     232.000us         7.25%     860.000us      53.750us          40 b           0 b            16            --
                                                 aten::linear         1.40%     166.000us         6.55%     777.000us     388.500us         800 b           0 b             2            --
                                       aten::l1_loss_backward         1.58%     188.000us         6.08%     721.000us     360.500us         400 b          -4 b             2            --
                                               aten::_to_copy         4.28%     508.000us         5.29%     628.000us      62.800us          40 b           0 b            10            --
                                                   aten::mean         0.89%     105.000us         4.44%     527.000us     263.500us           8 b           8 b             2            --
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 11.864ms
      """
      # 6-batch inference
      brain_analyst.evaluate(test_set=test_set)
      assert len(brain_analyst.profiler.speechbrain_event_traces) == 2
      # assert len(brain_analyst.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
      # assert len(brain_analyst.profiler.speechbrain_event_traces[1]) >= 125  # 144 with torch==1.10.1
      # as of evaluate() call
      # assert len(brain_analyst.profiler.events()) >= 125  # 142 with torch==1.10.1
      # assert len(brain_analyst.profiler.events().key_averages()) >= 25  # 42 with torch==1.10.1
      """print(brain_analyst.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls   Total FLOPs
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                ProfilerStep*        19.24%       1.129ms        96.92%       5.687ms       2.844ms         796 b      -1.61 Kb             2            --
                                                aten::l1_loss         5.16%     303.000us        35.50%       2.083ms     520.750us          16 b        -800 b             4            --
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         7.41%     435.000us        25.95%       1.523ms     761.500us       1.55 Kb         -40 b             2            --
                                                  aten::stack         6.15%     361.000us        18.20%       1.068ms     267.000us       1.56 Kb           0 b             4            --
                                                 aten::linear         3.43%     201.000us        15.78%     926.000us     463.000us         800 b           0 b             2            --
                                                   aten::mean         2.42%     142.000us        11.93%     700.000us     350.000us           8 b           8 b             2            --
                                               aten::isfinite         3.72%     218.000us        10.84%     636.000us     318.000us           2 b         -12 b             2            --
                                                    aten::cat         2.68%     157.000us         8.95%     525.000us     131.250us       1.56 Kb           0 b             4            --
                                                 aten::matmul         3.83%     225.000us         8.88%     521.000us     260.500us         800 b           0 b             2            --
                                                   aten::div_         3.34%     196.000us         6.97%     409.000us     204.500us           0 b          -8 b             2            --
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 5.868ms
      """
  
      brain_optimiser.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
      # key_avg_fit = deepcopy(brain_optimiser.profiler.events().key_averages())
      """print(brain_optimiser.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                ProfilerStep*        50.73%       1.874ms        98.86%       3.652ms       1.826ms         796 b      -1.66 Kb             2
                                                aten::l1_loss         1.38%      51.000us        11.83%     437.000us     109.250us          16 b        -400 b             4
                                               aten::isfinite         1.49%      55.000us         8.69%     321.000us     107.000us           3 b         -16 b             3
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         4.03%     149.000us         7.58%     280.000us      70.000us       1.55 Kb         -64 b             4
                                                 aten::linear         0.51%      19.000us         5.66%     209.000us     104.500us         800 b           0 b             2
                                                    aten::abs         3.19%     118.000us         5.58%     206.000us      25.750us          24 b          12 b             8
                                       aten::l1_loss_backward         0.92%      34.000us         4.28%     158.000us      79.000us         400 b          -4 b             2
                                                  aten::stack         1.00%      37.000us         3.87%     143.000us      28.600us       1.57 Kb           0 b             5
                                                  aten::empty         3.76%     139.000us         3.76%     139.000us       4.793us         544 b         544 b            29
                                                     aten::to         0.76%      28.000us         3.76%     139.000us       8.688us          44 b           4 b            16
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 3.694ms
      # to 11.864ms (analyst)
      """
  
      brain_optimiser.evaluate(test_set=test_set)
      """
      key_avg_evaluate = deepcopy(
          brain_optimiser.profiler.events().key_averages()
      )
      """
      """print(brain_optimiser.profiler.events().key_averages().table(sort_by="cpu_time_total", row_limit=10))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                ProfilerStep*        24.80%     524.000us        96.50%       2.039ms       1.020ms         796 b      -1.61 Kb             2
                                                aten::l1_loss         2.74%      58.000us        33.65%     711.000us     177.750us          16 b        -800 b             4
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         9.94%     210.000us        21.11%     446.000us     223.000us       1.55 Kb         -40 b             2
                                               aten::isfinite         3.64%      77.000us        15.76%     333.000us     166.500us           2 b         -12 b             2
                                                 aten::linear         1.04%      22.000us        11.88%     251.000us     125.500us         800 b           0 b             2
                                                   aten::mean         2.04%      43.000us        11.74%     248.000us     124.000us           8 b           8 b             2
                                                  aten::stack         3.31%      70.000us        10.18%     215.000us      53.750us       1.56 Kb           0 b             4
                                                   aten::div_         4.83%     102.000us         7.90%     167.000us      83.500us           0 b          -8 b             2
                                                 aten::matmul         1.61%      34.000us         7.38%     156.000us      78.000us         800 b           0 b             2
                                                     aten::ne         4.54%      96.000us         6.72%     142.000us      71.000us           2 b          -6 b             2
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 2.113ms
      # to 5.868ms (analyst)
      """
      # same check as for analyst
      assert len(brain_optimiser.profiler.speechbrain_event_traces) == 2
      # assert len(brain_optimiser.profiler.speechbrain_event_traces[0]) >= 250  # 296 with torch==1.10.1
      # assert len(brain_optimiser.profiler.speechbrain_event_traces[1]) >= 125  # 144 with torch==1.10.1
      # as of evaluate() call
      # assert len(brain_optimiser.profiler.events()) >= 125  # 142 with torch==1.10.1
      # assert len(brain_optimiser.profiler.events().key_averages()) >= 25  # 42 with torch==1.10.1
      # different config
      assert (
          brain_optimiser.profiler.record_shapes
          != brain_analyst.profiler.record_shapes
      )
      assert (
          brain_optimiser.profiler.with_stack
          != brain_analyst.profiler.with_stack
      )
      assert (
          brain_optimiser.profiler.with_flops
          != brain_analyst.profiler.with_flops
      )
      # same config
      assert (
          brain_optimiser.profiler.with_modules
          == brain_analyst.profiler.with_modules
      )
      assert (
          brain_optimiser.profiler.profile_memory
          == brain_analyst.profiler.profile_memory
      )
  
      """
      # let's take a look at the diff
      diff_fit, diff_evaluate = events_diff(key_avg_fit, key_avg_evaluate)
      # assert len(diff_fit) >= 50  # 64 with torch==1.10.1
      # assert len(diff_evaluate) >= 25  # 33 with torch==1.10.1
      # assert diff_fit.total_average().count >= 250  # 273 with torch==1.10.1
      # assert diff_evaluate.total_average().count >= 100  # 122 with torch==1.10.1
      " ""For curiosity only... the printed FunctionEvents differ by (name, # of Calls)
      print(diff_fit.table(sort_by="cpu_time_total"))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                               aten::isfinite         3.35%      55.000us        19.55%     321.000us     107.000us           3 b         -16 b             3
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...         9.07%     149.000us        17.05%     280.000us      70.000us       1.55 Kb         -64 b             4
                                                    aten::abs         7.19%     118.000us        12.55%     206.000us      25.750us          24 b          12 b             8
                                       aten::l1_loss_backward         2.07%      34.000us         9.62%     158.000us      79.000us         400 b          -4 b             2
                                                  aten::stack         2.25%      37.000us         8.71%     143.000us      28.600us       1.57 Kb           0 b             5
                                                  aten::empty         8.47%     139.000us         8.47%     139.000us       4.793us         544 b         544 b            29
                                                     aten::to         1.71%      28.000us         8.47%     139.000us       8.688us          44 b           4 b            16
                                               aten::_to_copy         1.89%      31.000us         7.00%     115.000us      11.500us          40 b           0 b            10
                                                     aten::mm         6.46%     106.000us         6.58%     108.000us      36.000us       1.17 Kb       1.17 Kb             3
                                                     aten::ne         4.26%      70.000us         6.46%     106.000us      35.333us           3 b          -9 b             3
                                                  aten::zeros         1.95%      32.000us         6.03%      99.000us      12.375us          32 b           0 b             8
      autograd::engine::evaluate_function: L1LossBackward0...         0.37%       6.000us         5.97%      98.000us      98.000us         396 b          -4 b             1
                                              L1LossBackward0         0.24%       4.000us         5.60%      92.000us      92.000us         400 b           0 b             1
                                                    aten::cat         0.97%      16.000us         5.05%      83.000us      16.600us       1.57 Kb           0 b             5
                                                    aten::div         2.80%      46.000us         4.81%      79.000us      26.333us          12 b           0 b             3
                                      Optimizer.step#SGD.step         3.71%      61.000us         4.51%      74.000us      74.000us          -4 b         -20 b             1
                                                   aten::_cat         2.13%      35.000us         4.08%      67.000us      13.400us       1.57 Kb           0 b             5
                                          aten::empty_strided         3.41%      56.000us         3.41%      56.000us       5.091us          44 b          44 b            11
             autograd::engine::evaluate_function: MmBackward0         0.37%       6.000us         3.41%      56.000us      56.000us           0 b        -400 b             1
                                                      aten::t         2.13%      35.000us         3.11%      51.000us      10.200us           0 b           0 b             5
                                                   aten::add_         3.05%      50.000us         3.05%      50.000us       8.333us           0 b           0 b             6
                                                  MmBackward0         0.43%       7.000us         3.05%      50.000us      50.000us         400 b           0 b             1
                                                    aten::mul         2.01%      33.000us         2.62%      43.000us      10.750us           5 b          -3 b             4
                                                    aten::sum         2.01%      33.000us         2.50%      41.000us      13.667us          40 b          40 b             3
            autograd::engine::evaluate_function: DivBackward0         0.55%       9.000us         2.38%      39.000us      39.000us          -4 b          -8 b             1
                                                   aten::norm         2.13%      35.000us         2.38%      39.000us      13.000us          12 b          12 b             3
                            Optimizer.zero_grad#SGD.zero_grad         1.46%      24.000us         1.89%      31.000us      31.000us          -4 b          -4 b             1
                                                 DivBackward0         0.24%       4.000us         1.83%      30.000us      30.000us           4 b           0 b             1
                                                   aten::item         1.04%      17.000us         1.77%      29.000us       4.143us           0 b           0 b             7
            autograd::engine::evaluate_function: AddBackward0         0.43%       7.000us         1.77%      29.000us      29.000us          40 b           0 b             1
                                                aten::random_         1.71%      28.000us         1.71%      28.000us      14.000us           0 b           0 b             2
                                                    aten::sub         1.71%      28.000us         1.71%      28.000us       9.333us         800 b         800 b             3
                                                  aten::copy_         1.71%      28.000us         1.71%      28.000us       2.800us           0 b           0 b            10
      autograd::engine::evaluate_function: torch::autograd...         0.30%       5.000us         1.71%      28.000us      14.000us        -440 b           0 b             2
                                                aten::resize_         1.64%      27.000us         1.64%      27.000us       5.400us       1.57 Kb       1.57 Kb             5
                                                    aten::add         0.85%      14.000us         1.64%      27.000us      27.000us           4 b           0 b             1
                                                     aten::eq         1.58%      26.000us         1.58%      26.000us       8.667us           3 b           3 b             3
                                              aten::unsqueeze         0.91%      15.000us         1.40%      23.000us       3.833us           0 b           0 b             6
                              torch::autograd::AccumulateGrad         0.97%      16.000us         1.40%      23.000us      11.500us        -440 b        -440 b             2
                                             aten::is_nonzero         0.43%       7.000us         1.40%      23.000us       7.667us           0 b           0 b             3
                                                 aten::detach         0.55%       9.000us         1.28%      21.000us       3.500us           0 b           0 b             6
                                             aten::as_strided         1.16%      19.000us         1.16%      19.000us       1.188us           0 b           0 b            16
                                                   aten::view         1.04%      17.000us         1.04%      17.000us       3.400us           0 b           0 b             5
                                              aten::transpose         0.67%      11.000us         0.97%      16.000us       3.200us           0 b           0 b             5
                                             aten::empty_like         0.37%       6.000us         0.91%      15.000us       7.500us         404 b           0 b             2
                                                       detach         0.73%      12.000us         0.85%      14.000us       2.333us           0 b           0 b             6
                                                  aten::clamp         0.61%      10.000us         0.85%      14.000us      14.000us           4 b           4 b             1
                                    aten::_local_scalar_dense         0.73%      12.000us         0.79%      13.000us       1.857us           0 b           0 b             7
                                                   aten::mul_         0.79%      13.000us         0.79%      13.000us       4.333us           0 b           0 b             3
                                              aten::ones_like         0.24%       4.000us         0.73%      12.000us      12.000us           4 b           0 b             1
                                             aten::zeros_like         0.30%       5.000us         0.73%      12.000us      12.000us         400 b           0 b             1
      autograd::engine::evaluate_function: UnsafeViewBackw...         0.18%       3.000us         0.73%      12.000us      12.000us           0 b           0 b             1
                                          UnsafeViewBackward0         0.12%       2.000us         0.55%       9.000us       9.000us           0 b           0 b             1
                                             aten::reciprocal         0.49%       8.000us         0.49%       8.000us       8.000us           4 b           4 b             1
                                                aten::reshape         0.24%       4.000us         0.43%       7.000us       7.000us           0 b           0 b             1
              autograd::engine::evaluate_function: TBackward0         0.12%       2.000us         0.43%       7.000us       7.000us           0 b           0 b             1
                                                  aten::zero_         0.37%       6.000us         0.37%       6.000us       0.545us           0 b           0 b            11
                                                  aten::fill_         0.37%       6.000us         0.37%       6.000us       1.500us           0 b           0 b             4
                                                   aten::sgn_         0.30%       5.000us         0.30%       5.000us       5.000us           0 b           0 b             1
                                                   TBackward0         0.06%       1.000us         0.30%       5.000us       5.000us           0 b           0 b             1
                                              aten::clamp_max         0.24%       4.000us         0.24%       4.000us       4.000us           0 b           0 b             1
                                         aten::_reshape_alias         0.18%       3.000us         0.18%       3.000us       3.000us           0 b           0 b             1
                                           aten::resolve_conj         0.12%       2.000us         0.12%       2.000us       0.333us           0 b           0 b             6
                                                 AddBackward0         0.06%       1.000us         0.06%       1.000us       1.000us           0 b           0 b             1
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 1.642ms
  
  
      print(diff_evaluate.table(sort_by="cpu_time_total"))
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
                                                         Name    Self CPU %      Self CPU   CPU total %     CPU total  CPU time avg       CPU Mem  Self CPU Mem    # of Calls
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      enumerate(DataLoader)#_SingleProcessDataLoaderIter._...        16.03%     210.000us        34.05%     446.000us     223.000us       1.55 Kb         -40 b             2
                                               aten::isfinite         5.88%      77.000us        25.42%     333.000us     166.500us           2 b         -12 b             2
                                                  aten::stack         5.34%      70.000us        16.41%     215.000us      53.750us       1.56 Kb           0 b             4
                                                     aten::ne         7.33%      96.000us        10.84%     142.000us      71.000us           2 b          -6 b             2
                                                  aten::empty        10.08%     132.000us        10.08%     132.000us       8.250us          80 b          80 b            16
                                                     aten::to         1.37%      18.000us         8.63%     113.000us      18.833us          16 b           0 b             6
                                                  aten::zeros         3.13%      41.000us         8.55%     112.000us      28.000us          16 b           0 b             4
                                                    aten::cat         1.45%      19.000us         8.40%     110.000us      27.500us       1.56 Kb           0 b             4
                                                     aten::mm         7.02%      92.000us         7.33%      96.000us      48.000us         800 b         800 b             2
                                                    aten::abs         4.58%      60.000us         7.33%      96.000us      16.000us          16 b           8 b             6
                                               aten::_to_copy         2.21%      29.000us         7.25%      95.000us      23.750us          16 b           0 b             4
                                                   aten::_cat         2.82%      37.000us         6.95%      91.000us      22.750us       1.56 Kb           0 b             4
                                                aten::resize_         3.51%      46.000us         3.51%      46.000us      11.500us       1.56 Kb       1.56 Kb             4
                                          aten::empty_strided         3.36%      44.000us         3.36%      44.000us      11.000us          16 b          16 b             4
                                                    aten::sub         2.98%      39.000us         2.98%      39.000us      19.500us         800 b         800 b             2
                                                    aten::sum         2.14%      28.000us         2.90%      38.000us      19.000us           0 b           0 b             2
                                                   aten::add_         2.82%      37.000us         2.82%      37.000us      18.500us           0 b           0 b             2
                                                      aten::t         1.68%      22.000us         2.75%      36.000us      18.000us           0 b           0 b             2
                                              aten::unsqueeze         1.98%      26.000us         2.67%      35.000us       8.750us           0 b           0 b             4
                                                     aten::eq         2.37%      31.000us         2.37%      31.000us      15.500us           2 b           2 b             2
                                                    aten::mul         2.37%      31.000us         2.37%      31.000us      15.500us           2 b           2 b             2
                                             aten::is_nonzero         0.53%       7.000us         1.76%      23.000us      11.500us           0 b           0 b             2
                                                   aten::item         0.92%      12.000us         1.76%      23.000us       5.750us           0 b           0 b             4
                                                  aten::copy_         1.68%      22.000us         1.68%      22.000us       5.500us           0 b           0 b             4
                                             aten::as_strided         1.37%      18.000us         1.37%      18.000us       2.250us           0 b           0 b             8
                                                   aten::view         1.30%      17.000us         1.30%      17.000us       4.250us           0 b           0 b             4
                                                 aten::detach         0.31%       4.000us         1.15%      15.000us       7.500us           0 b           0 b             2
                                              aten::transpose         0.69%       9.000us         1.07%      14.000us       7.000us           0 b           0 b             2
                                                       detach         0.84%      11.000us         0.84%      11.000us       5.500us           0 b           0 b             2
                                    aten::_local_scalar_dense         0.84%      11.000us         0.84%      11.000us       2.750us           0 b           0 b             4
                                                  aten::fill_         0.46%       6.000us         0.46%       6.000us       3.000us           0 b           0 b             2
                                                  aten::zero_         0.31%       4.000us         0.31%       4.000us       1.000us           0 b           0 b             4
                                           aten::resolve_conj         0.31%       4.000us         0.31%       4.000us       1.000us           0 b           0 b             4
      -------------------------------------------------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------  ------------
      Self CPU time total: 1.310ms
      """
  
      # set hook afterwards
      brain_analyst_raw = SimpleBrain(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      brain_optimiser_raw = SimpleBrain(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      brain_analyst_raw.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
      profile_analyst(brain_analyst_raw)
      brain_analyst_raw.evaluate(test_set=test_set)
      assert (
          getattr(brain_analyst_raw.profiler, "record_shapes") is True
      )
      assert getattr(brain_analyst_raw.profiler, "with_stack") is True
      assert getattr(brain_analyst_raw.profiler, "with_flops") is True
  
      brain_optimiser_raw.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
      profile_optimiser(brain_optimiser_raw)
      brain_optimiser_raw.evaluate(test_set=test_set)
      assert (
          getattr(brain_optimiser_raw.profiler, "record_shapes")
          is False
      )
      assert (
          getattr(brain_optimiser_raw.profiler, "with_stack") is False
      )
      assert (
          getattr(brain_optimiser_raw.profiler, "with_flops") is False
      )
  
      # wrap functions
      @profile_analyst
      def train_analyst(brain: SimpleBrain):
          brain.fit(
              epoch_counter=range(10),
              train_set=train_set,
              valid_set=valid_set,
          )
  
      @export
      @profile_optimiser
      def evaluate_optimiser(brain: SimpleBrain):
          brain.evaluate(test_set=test_set)
  
      brain_raw = SimpleBrain(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      assert brain_raw.profiler is None
      train_analyst(brain_raw)
      assert brain_raw.profiler is None
      evaluate_optimiser(brain_raw)
      assert brain_raw.profiler is None
  
      # profile classes
      @export
      @profile_analyst
      class SimpleBrainAnalyst(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])
  
      @profile_optimiser
      class SimpleBrainOptimiser(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])
  
      simple_brain_analyst = SimpleBrainAnalyst(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      assert (
          getattr(simple_brain_analyst.profiler, "record_shapes")
          is True
      )
      assert (
          getattr(simple_brain_analyst.profiler, "with_stack") is True
      )
      assert (
          getattr(simple_brain_analyst.profiler, "with_flops") is True
      )
      simple_brain_analyst.evaluate(test_set=test_set)
  
      simple_brain_optimiser = SimpleBrainOptimiser(
          {"model": model},
          lambda x: SGD(x, 0.1),
          run_opts={"device": device},
      )
      assert (
          getattr(simple_brain_optimiser.profiler, "record_shapes")
          is False
      )
      assert (
          getattr(simple_brain_optimiser.profiler, "with_stack")
          is False
      )
      assert (
          getattr(simple_brain_optimiser.profiler, "with_flops")
          is False
      )
      simple_brain_optimiser.fit(
          epoch_counter=range(10),
          train_set=train_set,
          valid_set=valid_set,
      )
  
label_desc: Use image node_id 6 for calling a global function with example usage: test_profile_details(device) and returns: self, torch, self, torch, self, torch