Update evaluator.py

microsoft · Jun 12, 2023 · 00b9843 · 00b9843
1 parent cfb8784
commit 00b9843
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nni/contrib/compression/utils/evaluator.py b/nni/contrib/compression/utils/evaluator.py
@@ -394,8 +394,8 @@ def __init__(self, trainer: pl.Trainer, data_module: pl.LightningDataModule | L.
         )
         lighting2_check = not LIGHTNING2_INSTALLED and isinstance(data_module, L.LightningDataModule)
         assert isinstance(data_module, pl.LightningDataModule) and is_traceable(data_module), err_msg
-        self.trainer = trainer
-        self.data_module = data_module
+        self.trainer: pl.Trainer | L.Trainer = trainer
+        self.data_module: pl.LightningDataModule | L.LightningDataModule = data_module
         self._dummy_input = dummy_input
 
         self.model: pl.LightningModule | L.LightningModule | None = None