[Fix] Fix test

mmengine/model/base_model/base_model.py

@@ -116,7 +116,7 @@ def train_step(self, data: Union[dict, tuple, list],
         optim_wrapper.update_params(parsed_losses)
         return log_vars
 
-    def val_step(self, data: Union[tuple, dict, list]) -> Union[tuple, list]:
+    def val_step(self, data: Union[tuple, dict, list]) -> list:
         """Gets the predictions of given data.
 
         Calls ``self.data_preprocessor(data, False)`` and
@@ -132,7 +132,7 @@ def val_step(self, data: Union[tuple, dict, list]) -> Union[tuple, list]:
         data = self.data_preprocessor(data, False)
         return self._run_forward(data, mode='predict')  # type: ignore
 
-    def test_step(self, data: Union[dict, tuple, list]) -> Union[tuple, list]:
+    def test_step(self, data: Union[dict, tuple, list]) -> list:
         """``BaseModel`` implements ``test_step`` the same as ``val_step``.
 
         Args:

mmengine/runner/loops.py

@@ -10,6 +10,7 @@
 from mmengine.evaluator import Evaluator
 from mmengine.logging import print_log
 from mmengine.registry import LOOPS
+from mmengine.structures import BaseDataElement
 from mmengine.utils import is_list_of
 from .amp import autocast
 from .base_loop import BaseLoop
@@ -399,12 +400,13 @@ def run_iter(self, idx, data_batch: Sequence[dict]):
             'before_val_iter', batch_idx=idx, data_batch=data_batch)
         # outputs should be sequence of BaseDataElement
         with autocast(enabled=self.fp16):
-            results = self.runner.model.test_step(data_batch)
-        outputs, loss = list(), dict()
-        if isinstance(results, tuple):
-            outputs, loss = results
-        elif isinstance(results, list):
-            outputs, loss = results, dict()
+            outputs = self.runner.model.test_step(data_batch)
+        if isinstance(outputs[-1],
+                      BaseDataElement) and outputs.keys() == ['loss']:
+            loss = outputs[-1].loss  # type: ignore
+            outputs = outputs[:-1]
+        else:
+            loss = dict()
         # get val loss and avoid breaking change
         for loss_name, loss_value in loss.items():
             if loss_name not in self.val_loss:
@@ -494,12 +496,13 @@ def run_iter(self, idx, data_batch: Sequence[dict]) -> None:
             'before_test_iter', batch_idx=idx, data_batch=data_batch)
         # predictions should be sequence of BaseDataElement
         with autocast(enabled=self.fp16):
-            results = self.runner.model.test_step(data_batch)
-        outputs, loss = list(), dict()
-        if isinstance(results, tuple):
-            outputs, loss = results
-        elif isinstance(results, list):
-            outputs, loss = results, dict()
+            outputs = self.runner.model.test_step(data_batch)
+        if isinstance(outputs[-1],
+                      BaseDataElement) and outputs.keys() == ['loss']:
+            loss = outputs[-1].loss  # type: ignore
+            outputs = outputs[:-1]
+        else:
+            loss = dict()
         # get val loss and avoid breaking change
         for loss_name, loss_value in loss.items():
             if loss_name not in self.test_loss: