class_docstr: Sampling-based super-net training but using an RL agent to control the sampling.
  
  The default implementation for :class:`~nni.nas.strategy.ENAS`.
  
  See Also
  --------
  nni.nas.strategy.ENAS
  RandomSamplingLightningModule
code: def training_step(self, batch_packed, batch_idx):
      if len(batch_packed) == 2:
          # Legacy (pytorch-lightning 1.x): The received batch is a tuple of (data, "train" | "val")
          batch, mode = batch_packed
      else:
          # New (pytorch-lightning 2.0+): a tuple of data, batch_idx, and dataloader_idx
          batch, _, dataloader_idx = batch_packed
          mode = "train" if dataloader_idx == 0 else "val"
  
      assert self._replay_buffer is not None
  
      step_output = None  # Sometimes step can be skipped
  
      if mode == "train":
          # train model params
          self.policy.eval()
          self.resample()
          step_output = self.training_module.training_step(
              batch, batch_idx
          )
          w_step_loss = (
              step_output["loss"]
              if isinstance(step_output, dict)
              else step_output
          )
          self.advance_optimization(w_step_loss, batch_idx)
  
      elif (
          self.warmup_epochs == 0
          or self.trainer.current_epoch >= self.warmup_epochs
      ):
          # Run a sample to retrieve the reward
          self.policy.train()
          sample = self.resample()
          step_output = self.training_module.validation_step(
              batch, batch_idx
          )
  
          # use the default metric of self.training_module as reward function
          if len(self.trainer.callback_metrics) == 1:
              _, metric = next(
                  iter(self.trainer.callback_metrics.items())
              )
          else:
              metric_name = self.reward_metric_name or "default"
              if metric_name not in self.trainer.callback_metrics:
                  raise KeyError(
                      f"Model reported metrics should contain a ``{metric_name}`` key but "
                      f"found multiple (or zero) metrics without default: {list(self.trainer.callback_metrics.keys())}. "
                      f"Please try to set ``reward_metric_name`` to be one of the keys listed above. "
                      f"If it is not working use self.log to report metrics with the specified key ``{metric_name}`` "
                      "in validation_step, and remember to set on_step=True."
                  )
              metric = self.trainer.callback_metrics[metric_name]
          reward: float = metric.item()
  
          if self.penalty is not None:
              if isinstance(self.penalty, SampleProfilerPenalty):
                  reward, details = self.penalty(reward, sample)
              elif isinstance(self.penalty, ExpectationProfilerPenalty):
                  reward, details = self.penalty(
                      reward, self.export_probs()
                  )
              else:
                  raise TypeError(
                      f"Unknown penalty type: {type(self.penalty)}"
                  )
              self.log_dict(
                  {f"penalty/{k}": v for k, v in details.items()}
              )
  
          # Put it into replay buffer
          trajectory = self.generator.send_reward(reward)
          self._replay_buffer.update(trajectory)
          self._trajectory_counter += 1
  
      self.advance_lr_schedulers(batch_idx)
  
      if (
          self._trajectory_counter > 0
          and self._trajectory_counter % self.log_prob_every_n_step == 0
      ):
          self.log_probs(self.export_probs())
  
      if (
          self._trajectory_counter > 0
          and self._trajectory_counter % self.batches_per_update == 0
      ):
          # Export could be just called.
          # The policy must be in train mode to make update work.
          self.policy.train()
          update_times = self.update_kwargs.get("update_times", 1)
          for _ in range(update_times):
              self.policy.update(
                  0, self._replay_buffer, **self.update_kwargs
              )
  
      return step_output
  
imports: from __future__ import annotations
  import warnings
  import logging
  from typing import Any, Callable, TYPE_CHECKING
  import pytorch_lightning
  import torch
  from nni.mutable import Sample
  from nni.nas.nn.pytorch import ModelSpace
  from .base_lightning import BaseOneShotLightningModule
  from .profiler import ProfilerPenalty, ExpectationProfilerPenalty, SampleProfilerPenalty
label_desc: Use image node_id 5 for calling the EnasLightningModule obj's underlying member method code with example usage: obj.training_step(batch_packed, batch_idx) and returns: step_output