[zero] added error message to handle on-the-fly import of torch Module class
diff --git a/colossalai/utils/model/utils.py b/colossalai/utils/model/utils.py
index 0b0b73820a85..ecc0cdb5a1dd 100644
--- a/colossalai/utils/model/utils.py
+++ b/colossalai/utils/model/utils.py
@@ -80,6 +80,10 @@ def __exit__(self, exc_type, exc_value, traceback):
             torch.set_default_dtype(self._old_default_dtype)
 
         def _disable_class(cls):
+            if not hasattr(cls, '_old_init'):
+                raise AttributeError(
+                    f"_old_init is not found in the {cls.__name__}, please make sure that you have imported {cls.__name__} before entering the context."
+                )
             cls.__init__ = cls._old_init
 
         # Replace .__init__() for all existing subclasses of torch.nn.Module

The zero init context has a precondition that the module class has to be imported before entering the context. However, sometimes this is not always so. For example:

```
with ctx:
    model = transformers.BertModel(...)
```

This example will raise exception because BertModel is only imported on the fly. This PR adds an error message for easier debugging.