[context] maintain the context object in with statement (#1073)

bfdc5ccb · Frank Lee · GitHub · c5cd3b0f · bfdc5ccb
Unverified Commit bfdc5ccb authored Jun 07, 2022 by Frank Lee Committed by GitHub Jun 07, 2022
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

colossalai/utils/model/utils.py colossalai/utils/model/utils.py +1 -0

No files found.
--- a/colossalai/utils/model/utils.py
+++ b/colossalai/utils/model/utils.py
@@ -72,6 +72,7 @@ class InsertPostInitMethodToModuleSubClasses(object):
        torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass)

        self._pre_context_exec()
+        return self

    def __exit__(self, exc_type, exc_value, traceback):