class_docstr: Extend :class:`DartsLightningModule` to support gumbel-softmax with temperature annealing.
  
  The default implementation of :class:`~nni.nas.strategy.GumbelDARTS`.
  
  See Also
  --------
  nni.nas.strategy.GumbelDARTS
  DartsLightningModule
class_docstr_tok: ["Extend",":","class",":","`","DartsLightningModule","`","to","support","gumbel-softmax","with","temperature","annealing",".","The","default","implementation","of",":","class",":","`","~nni.nas.strategy.GumbelDARTS","`",".","See","Also","--","--","--","--","nni.nas.strategy.GumbelDARTS","DartsLightningModule"]
code: def on_train_epoch_start(self):
      temp = self.temperature_scheduler.step(
          self.trainer.current_epoch, self.trainer.max_epochs
      )
      self.log("gumbel_temperature", temp)
  
      for module in self.modules():
          if isinstance(module, GumbelSoftmax):
              module.tau = temp  # type: ignore
  
      return super().on_train_epoch_start()
  
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 GumbelDartsLightningModule obj's underlying member method code with example usage: obj.on_train_epoch_start() and returns: super