minor change

open-mmlab · Aug 26, 2022 · d676d93 · d676d93
1 parent 2fbdd01
commit d676d93
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/mmrazor/models/algorithms/nas/dsnas.py b/mmrazor/models/algorithms/nas/dsnas.py
@@ -177,7 +177,7 @@ def train_step(self, data: List[dict],
             # 2. update mutator
             if self.update_mutator(cur_epoch):
                 with optim_wrapper['mutator'].optim_context(self):
-                    mutator_loss = self.compute_mutator_loss(cur_epoch)
+                    mutator_loss = self.compute_mutator_loss()
                 mutator_losses, mutator_log_vars = \
                     self.parse_losses(mutator_loss)
                 optim_wrapper['mutator'].update_params(mutator_losses)
@@ -216,15 +216,12 @@ def update_mutator(self, cur_epoch: int) -> bool:
             return True
         return False
 
-    def compute_mutator_loss(self, cur_epoch: int) -> Dict[str, torch.Tensor]:
+    def compute_mutator_loss(self) -> Dict[str, torch.Tensor]:
         """Compute mutator loss.
 
         In this method, arch_loss & flops_loss[optional] are computed
         by traversing arch_weights & probs in search groups.
 
-        Args:
-            cur_epoch (int): Current training epoch.
-
         Returns:
             Dict: Loss of the mutator.
         """
@@ -317,7 +314,7 @@ def train_step(self, data: List[dict],
             # 2. update mutator
             if self.module.update_mutator(cur_epoch):
                 with optim_wrapper['mutator'].optim_context(self):
-                    mutator_loss = self.module.compute_mutator_loss(cur_epoch)
+                    mutator_loss = self.module.compute_mutator_loss()
                 mutator_losses, mutator_log_vars = \
                     self.module.parse_losses(mutator_loss)
                 optim_wrapper['mutator'].update_params(mutator_losses)

diff --git a/mmrazor/models/architectures/generators/base_generator.py b/mmrazor/models/architectures/generators/base_generator.py
@@ -2,7 +2,7 @@
 from typing import Dict, Optional
 
 import torch
-from mmengine.runner import BaseModule
+from mmengine.model import BaseModule
 
 from mmrazor.models.utils import get_module_device