Fix typecheck for pytorch lightning v1.8 (#5207)

nni/nas/oneshot/pytorch/base_lightning.py

@@ -388,7 +388,7 @@ class BaseOneShotLightningModule(pl.LightningModule):
 
         return optim_conf
 
-    def setup(self, stage=None):
+    def setup(self, stage: str = cast(str, None)):  # add default value to be backward-compatible
         # redirect the access to trainer/log to this module
         # but note that we might be missing other attributes,
         # which could potentially be a problem
@@ -400,7 +400,7 @@ class BaseOneShotLightningModule(pl.LightningModule):
 
         return self.model.setup(stage)
 
-    def teardown(self, stage=None):
+    def teardown(self, stage: str = cast(str, None)):
         return self.model.teardown(stage)
 
     def configure_architecture_optimizers(self) -> list[optim.Optimizer] | optim.Optimizer | None:
@@ -492,7 +492,7 @@ class BaseOneShotLightningModule(pl.LightningModule):
                     self.model.lr_scheduler_step(cast(Any, scheduler), cast(int, opt_idx), None)
         except AttributeError:
             # lightning < 1.6
-            for lr_scheduler in self.trainer.lr_schedulers:
+            for lr_scheduler in self.trainer.lr_schedulers:  # type: ignore
                 if lr_scheduler['reduce_on_plateau']:
                     warnings.warn('Reduce-lr-on-plateau is not supported in NAS. It will be ignored.', UserWarning)
                 if lr_scheduler['interval'] == interval and current_idx % lr_scheduler['frequency']:

nni/nas/oneshot/pytorch/differentiable.py

@@ -275,6 +275,8 @@ class GumbelDartsLightningModule(DartsLightningModule):
 
     def on_train_epoch_start(self):
         if self.use_temp_anneal:
+            if self.trainer.max_epochs is None:
+                raise ValueError('Please set max_epochs for trainer when using temperature annealing.')
             self.temp = (1 - self.trainer.current_epoch / self.trainer.max_epochs) * (self.init_temp - self.min_temp) + self.min_temp
             self.temp = max(self.temp, self.min_temp)