[Fix] Fix base tta model (#593)

Co-authored-by: ubuntu <ubuntu@localhost.localdomain>
open-mmlab · Oct 11, 2022 · 2df5bc1 · 2df5bc1
1 parent 46add35
commit 2df5bc1
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/mmengine/model/wrappers/test_time_aug.py b/mmengine/model/wrappers/test_time_aug.py
@@ -20,7 +20,7 @@
 
 
 @MODELS.register_module()
-class BaseTTAModel:
+class BaseTTAModel(nn.Module):
     """Base model for inference with test-time augmentation.
 
     ``BaseTTAModel`` is a wrapper for inference given multi-batch data.
@@ -74,6 +74,7 @@ class BaseTTAModel:
     """
 
     def __init__(self, module: Union[dict, nn.Module]):
+        super().__init__()
         if isinstance(module, nn.Module):
             self.module = module
         elif isinstance(module, dict):