class_docstr: Search implementation of :class:`~nni.nas.strategy.DARTS`.
  
  The dataloader is expected to be a CombinedLoader with two keys: ``train`` and ``val``.
  
  See Also
  --------
  nni.nas.strategy.DARTS
  nni.nas.pytorch.oneshot.base_lightning.BaseOneShotLightningModule
class_docstr_tok: ["Search","implementation","of",":","class",":","`","~nni.nas.strategy.DARTS","`",".","The","dataloader","is","expected","to","be","a","CombinedLoader","with","two","keys",":","``","train","``","and","``","val","``",".","See","Also","--","--","--","--","nni.nas.strategy.DARTS","nni.nas.pytorch.oneshot.base_lightning.BaseOneShotLightningModule"]
code: def training_step(self, batch, batch_idx):
      # grad manually
      arc_optim = self.architecture_optimizers()
      if not isinstance(arc_optim, optim.Optimizer):
          raise TypeError(
              f"Expect arc_optim to be a single Optimizer, but found: {arc_optim}"
          )
  
      # DARTS strategy makes sure that ``train`` and ``val`` must be in the batch
      trn_batch = batch["train"]
      val_batch = batch["val"]
  
      # Phase 1: architecture step. Only when warmup completes.
      if (
          self.warmup_epochs == 0
          or self.trainer.current_epoch >= self.warmup_epochs
      ):
          # The _resample hook is kept for some darts-based NAS methods like proxyless.
          # See code of those methods for details.
          self.resample()
          arc_optim.zero_grad()
          arc_step_loss = self.training_module.training_step(
              val_batch, 2 * batch_idx
          )
          if isinstance(arc_step_loss, dict):
              arc_step_loss = arc_step_loss["loss"]
  
          if self.penalty is not None:
              arc_step_loss, details = self.penalty(
                  arc_step_loss, self.export_probs()
              )
              self.log_dict(
                  {f"penalty/{k}": v for k, v in details.items()}
              )
  
          self.manual_backward(arc_step_loss)
          arc_optim.step()
  
      # Phase 2: model step
      self.resample()
      loss_and_metrics = self.training_module.training_step(
          trn_batch, 2 * batch_idx + 1
      )
      w_step_loss = (
          loss_and_metrics["loss"]
          if isinstance(loss_and_metrics, dict)
          else loss_and_metrics
      )
      self.advance_optimization(
          w_step_loss, batch_idx, self.gradient_clip_val
      )
  
      # Update learning rates
      self.advance_lr_schedulers(batch_idx)
  
      if batch_idx % self.log_prob_every_n_step == 0:
          # NOTE: Ideally we should use global_step, but, who cares.
          self.log_probs(self.export_probs())
  
      return loss_and_metrics
  
imports: from __future__ import annotations
  import logging
  import pytorch_lightning
  import torch
  import torch.nn
  import torch.optim
  from .base_lightning import BaseOneShotLightningModule
  from .supermodule.differentiable import GumbelSoftmax
  from .profiler import ExpectationProfilerPenalty
label_desc: Use image node_id 2 for calling the DartsLightningModule obj's underlying member method code with example usage: obj.training_step(batch, batch_idx) and returns: loss_and_metrics