compile train_step, val_step, test_step instead

open-mmlab · Mar 8, 2023 · b845d98 · b845d98
1 parent 7b4060b
commit b845d98
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/mmengine/runner/runner.py b/mmengine/runner/runner.py
@@ -2327,8 +2327,11 @@ def _maybe_compile(self) -> None:
         assert isinstance(compile_cfg, dict), (
             f'`compile` should be a dict or bool, got {type(compile_cfg)}')
 
-        # Compile the model.forward
-        self.model.forward = torch.compile(self.model.forward, **compile_cfg)
+        # Compile the model's train_step, val_step and test_step
+        for target in ('train_step', 'val_step', 'test_step'):
+            func = getattr(self.model, target)
+            compiled_func = torch.compile(func, **compile_cfg)
+            setattr(self.model, target, compiled_func)
         self._is_compiled = True
         self.logger.info('Model has been "compiled". The first few iterations'
                          ' will be slow, please be patient.')