class_docstr: This class is used to override the schedule logic from the profiler and perform recording for both
  `training_step`, `validation_step`.
class_docstr_tok: ["This","class","is","used","to","override","the","schedule","logic","from","the","profiler","and","perform","recording","for","both","`","training_step","`",",","`","validation_step","`","."]
code: def __call__(self, num_step: int) -> "ProfilerAction":
      # ignore the provided input. Keep internal state instead.
      if self._current_action is None or self.has_finished:
          return ProfilerAction.NONE
  
      self._step()
      action = self._schedule(max(self.num_step, 0))
      if (
          self._prev_schedule_action == ProfilerAction.RECORD
          and action == ProfilerAction.WARMUP
      ):
          # Work around the corner case when validation starts before train.
          # In this case, the action is RECORD in validation loop, and then call into the train
          # and the action is still WARMUP in train and pytorch will recognize this as error.
          action = ProfilerAction.RECORD
      if action == ProfilerAction.RECORD_AND_SAVE:
          if self.is_training:
              self._training_step_reached_end = True
          elif self.is_validating:
              self._validation_step_reached_end = True
          elif self.is_testing:
              self._test_step_reached_end = True
          elif self.is_predicting:
              self._predict_step_reached_end = True
      self._prev_schedule_action = action
      return action
  
imports: import inspect
  import logging
  import os
  from functools import lru_cache, partial
  from pathlib import Path
  from typing import TYPE_CHECKING, Any, Callable, ContextManager, Dict, List, Optional, Type, Union
  import torch
  from torch import Tensor, nn
  from torch.autograd.profiler import EventList, record_function
  from torch.profiler import ProfilerAction, ProfilerActivity, tensorboard_trace_handler
  from torch.utils.hooks import RemovableHandle
  from typing_extensions import override
  from lightning.fabric.accelerators.cuda import is_cuda_available
  from lightning.pytorch.profilers.profiler import Profiler
  from lightning.pytorch.utilities.exceptions import MisconfigurationException
  from lightning.pytorch.utilities.rank_zero import WarningCache, rank_zero_warn
label_desc: Use image node_id 12 for calling the ScheduleWrapper obj's underlying member method code with example usage: obj.__call__(num_step) and returns: action, ProfilerAction