[context] maintain the context object in with statement

[FrankLeeeee]

The InsertPostInitMethodToModuleSubClasses context manager does not return the object in with statement, this issue fixed this problem so that the attributes of the contexts can still be accessed after with statement.

Previous:

>>> with InsertPostInitMethodToModuleSubClasses() as ctx:
>>>        model = resnet18()
>>> print(ctx)
None

Now:

>>> with InsertPostInitMethodToModuleSubClasses() as ctx:
>>>        model = resnet18()
>>> print(ctx)
<colossalai.utils.model.colo_init_context.ColoInitContext at 0x7f75d031bcd0>