From 007ca0df925350c70ee8ff324599f8e156cec59a Mon Sep 17 00:00:00 2001 From: ver217 Date: Wed, 25 May 2022 20:41:58 +0800 Subject: [PATCH] fix colo init context (#1026) --- colossalai/utils/model/colo_init_context.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/colossalai/utils/model/colo_init_context.py b/colossalai/utils/model/colo_init_context.py index 0eada7583..a39353b16 100644 --- a/colossalai/utils/model/colo_init_context.py +++ b/colossalai/utils/model/colo_init_context.py @@ -69,7 +69,7 @@ def _setattr_with_colotensor(self, name: str, value: Union[torch.Tensor, torch.n d.discard(name) params = self.__dict__.get('_parameters') - if isinstance(value, (ColoTensor, torch.nn.Parameter)): + if isinstance(value, (ColoParameter, torch.nn.Parameter)): if params is None: raise AttributeError("cannot assign parameters before Module.__init__() call") remove_from(self.__dict__, self._buffers, self._modules, self._non_persistent_buffers_set)