diff --git a/colossalai/nn/__init__.py b/colossalai/nn/__init__.py
index 3991e3bfb..7c37650ed 100644
--- a/colossalai/nn/__init__.py
+++ b/colossalai/nn/__init__.py
@@ -4,3 +4,7 @@ from .lr_scheduler import *
 from .metric import *
 from .model import *
 from .optimizer import *
+from ._ops import *
+
+from .modules import ColoLinear, ColoEmbedding
+from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module
diff --git a/colossalai/tensor/_ops/__init__.py b/colossalai/nn/_ops/__init__.py
similarity index 100%
rename from colossalai/tensor/_ops/__init__.py
rename to colossalai/nn/_ops/__init__.py
diff --git a/colossalai/tensor/_ops/_utils.py b/colossalai/nn/_ops/_utils.py
similarity index 100%
rename from colossalai/tensor/_ops/_utils.py
rename to colossalai/nn/_ops/_utils.py
diff --git a/colossalai/tensor/_ops/addmm.py b/colossalai/nn/_ops/addmm.py
similarity index 100%
rename from colossalai/tensor/_ops/addmm.py
rename to colossalai/nn/_ops/addmm.py
diff --git a/colossalai/tensor/_ops/element_wise.py b/colossalai/nn/_ops/element_wise.py
similarity index 100%
rename from colossalai/tensor/_ops/element_wise.py
rename to colossalai/nn/_ops/element_wise.py
diff --git a/colossalai/tensor/_ops/embedding.py b/colossalai/nn/_ops/embedding.py
similarity index 100%
rename from colossalai/tensor/_ops/embedding.py
rename to colossalai/nn/_ops/embedding.py
diff --git a/colossalai/tensor/_ops/layernorm.py b/colossalai/nn/_ops/layernorm.py
similarity index 100%
rename from colossalai/tensor/_ops/layernorm.py
rename to colossalai/nn/_ops/layernorm.py
diff --git a/colossalai/tensor/_ops/linear.py b/colossalai/nn/_ops/linear.py
similarity index 100%
rename from colossalai/tensor/_ops/linear.py
rename to colossalai/nn/_ops/linear.py
diff --git a/colossalai/tensor/_ops/loss.py b/colossalai/nn/_ops/loss.py
similarity index 100%
rename from colossalai/tensor/_ops/loss.py
rename to colossalai/nn/_ops/loss.py
diff --git a/colossalai/tensor/module_utils.py b/colossalai/nn/module_utils.py
similarity index 92%
rename from colossalai/tensor/module_utils.py
rename to colossalai/nn/module_utils.py
index 9fa389171..64d3f8075 100644
--- a/colossalai/tensor/module_utils.py
+++ b/colossalai/nn/module_utils.py
@@ -10,6 +10,7 @@ def register_colo_module(module_type: type, colo_module: ColoModule):
     global _COLOSSAL_MODULES
     _COLOSSAL_MODULES[module_type] = colo_module
 
+
 def is_colo_module(module: torch.nn.Module):
     global _COLOSSAL_MODULES
     for module_type in _COLOSSAL_MODULES.keys():
@@ -17,6 +18,7 @@ def is_colo_module(module: torch.nn.Module):
             return True
     return False
 
+
 def get_colo_module(module: torch.nn.Module):
     global _COLOSSAL_MODULES
     if is_colo_module(module):
@@ -26,6 +28,7 @@ def get_colo_module(module: torch.nn.Module):
     else:
         return None
 
+
 def check_colo_module(module: torch.nn.Module, recursive=True):
     if is_colo_module(module):
         colo_module = get_colo_module(module)
@@ -35,20 +38,22 @@ def check_colo_module(module: torch.nn.Module, recursive=True):
             param = module.get_parameter(param_name)
             if not isinstance(param, ColoParameter):
                 raise Exception(f'Invalid ColoParameter spec: {param} in {module} is not a ColoParameter.')
-            if param.has_spec():  
+            if param.has_spec():
                 cur_compute_pattern = param.spec.parallel_action.compute_pattern
                 if compute_pattern is None:
                     compute_pattern = cur_compute_pattern
                 else:
                     if cur_compute_pattern != compute_pattern:
-                        raise Exception(f'Invalid ColoParameter spec: Params in {module} have different compute_pattern.')
+                        raise Exception(
+                            f'Invalid ColoParameter spec: Params in {module} have different compute_pattern.')
             else:
                 continue
-            
+
         if compute_pattern is not None:
             colo_module.register(compute_pattern)
             if not colo_module.has_compute_pattern(compute_pattern):
-                raise Exception(f'Invalid ColoParameter spec: ComputePattern {compute_pattern} in {module} is not allowed.')
+                raise Exception(
+                    f'Invalid ColoParameter spec: ComputePattern {compute_pattern} in {module} is not allowed.')
 
             match_specs = False
             allowed_specs = colo_module.get_dist_specs(compute_pattern)
@@ -73,6 +78,7 @@ def check_colo_module(module: torch.nn.Module, recursive=True):
         for submodule in module.children():
             check_colo_module(submodule, recursive=True)
 
+
 def init_colo_module(module: torch.nn.Module, parallel_action: ParallelAction, recursive=True, mode='default'):
     compute_pattern = parallel_action.compute_pattern
     if is_colo_module(module):
@@ -99,4 +105,3 @@ def init_colo_module(module: torch.nn.Module, parallel_action: ParallelAction, r
     if recursive == True:
         for submodule in module.children():
             init_colo_module(submodule, parallel_action, recursive=True, mode=mode)
-    
\ No newline at end of file
diff --git a/colossalai/tensor/modules/__init__.py b/colossalai/nn/modules/__init__.py
similarity index 100%
rename from colossalai/tensor/modules/__init__.py
rename to colossalai/nn/modules/__init__.py
diff --git a/colossalai/tensor/modules/colo_module.py b/colossalai/nn/modules/colo_module.py
similarity index 81%
rename from colossalai/tensor/modules/colo_module.py
rename to colossalai/nn/modules/colo_module.py
index ecdfc1a59..2022411f9 100644
--- a/colossalai/tensor/modules/colo_module.py
+++ b/colossalai/nn/modules/colo_module.py
@@ -4,11 +4,12 @@ from typing import List, Dict
 
 
 class ColoModule(object):
+
     def __init__(self):
         self._shard_params: List[str] = []
         # Example:
         # {ComputePattern.TP1D:
-        #    'default': 
+        #    'default':
         #      'weight':
         #          distspec.shard(xxxxx)
         #      'bias':
@@ -21,25 +22,29 @@ class ColoModule(object):
     def _register_shard_params(self, params: List[str]):
         self._shard_params = params
 
-    def _register_allowed_patterns(self, compute_pattern: ComputePattern, dist_specs: Dict[str, _DistSpec], mode='default'):
-        assert list(dist_specs.keys()).sort() == self._shard_params.sort(), 'Every registered param should have dist_spec.'
+    def _register_allowed_patterns(self,
+                                   compute_pattern: ComputePattern,
+                                   dist_specs: Dict[str, _DistSpec],
+                                   mode='default'):
+        assert list(
+            dist_specs.keys()).sort() == self._shard_params.sort(), 'Every registered param should have dist_spec.'
         if not compute_pattern in self._allowed_patterns:
             self._allowed_patterns[compute_pattern] = {}
         self._allowed_patterns[compute_pattern][mode] = dist_specs
 
     def _set_default(self, compute_pattern: ComputePattern, target_mode):
         self._allowed_patterns[compute_pattern]['default'] = self._allowed_patterns[compute_pattern][target_mode]
-    
+
     def has_compute_pattern(self, compute_pattern: ComputePattern):
         return compute_pattern in self._allowed_patterns
-    
+
     def get_dist_specs(self, compute_pattern: ComputePattern):
         assert self.has_compute_pattern(compute_pattern)
         return self._allowed_patterns[compute_pattern]
-        
+
     def has_compute_pattern_with_mode(self, compute_pattern: ComputePattern, mode='default'):
         return compute_pattern in self._allowed_patterns and mode in self._allowed_patterns[compute_pattern]
-        
+
     def get_dist_specs_with_mode(self, compute_pattern: ComputePattern, mode='default'):
         assert self.has_compute_pattern_with_mode(compute_pattern, mode)
         return self._allowed_patterns[compute_pattern][mode]
@@ -48,4 +53,4 @@ class ColoModule(object):
         return self._shard_params
 
     def register(self, compute_pattern):
-        raise NotImplementedError
\ No newline at end of file
+        raise NotImplementedError
diff --git a/colossalai/tensor/modules/embedding.py b/colossalai/nn/modules/embedding.py
similarity index 72%
rename from colossalai/tensor/modules/embedding.py
rename to colossalai/nn/modules/embedding.py
index b48193971..9dd575a5b 100644
--- a/colossalai/tensor/modules/embedding.py
+++ b/colossalai/nn/modules/embedding.py
@@ -3,23 +3,27 @@ from colossalai.tensor import ComputePattern, distspec
 from colossalai.core import global_context as gpc
 from colossalai.context.parallel_mode import ParallelMode
 
+
 class ColoEmbedding(ColoModule):
+
     def __init__(self):
         super(ColoEmbedding, self).__init__()
         self._register_shard_params(['weight'])
-        
+
     def register(self, compute_pattern):
         if not compute_pattern in self._allowed_patterns:
             if ComputePattern.TP1D == compute_pattern:
                 self._set_TP1D()
-        
+
     def _set_TP1D(self):
         # TP1D Row Linear
         _compute_pattern = ComputePattern.TP1D
         self._register_allowed_patterns(
             compute_pattern=_compute_pattern,
             dist_specs={
-                'weight': distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
             },
             mode='row',
         )
@@ -28,9 +32,11 @@ class ColoEmbedding(ColoModule):
         self._register_allowed_patterns(
             compute_pattern=_compute_pattern,
             dist_specs={
-                'weight': distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
             },
             mode='col',
         )
 
-        self._set_default(compute_pattern=_compute_pattern, target_mode='row')
\ No newline at end of file
+        self._set_default(compute_pattern=_compute_pattern, target_mode='row')
diff --git a/colossalai/tensor/modules/linear.py b/colossalai/nn/modules/linear.py
similarity index 63%
rename from colossalai/tensor/modules/linear.py
rename to colossalai/nn/modules/linear.py
index 239cbeb1e..1cf32daf2 100644
--- a/colossalai/tensor/modules/linear.py
+++ b/colossalai/nn/modules/linear.py
@@ -3,24 +3,29 @@ from colossalai.tensor import ComputePattern, distspec
 from colossalai.core import global_context as gpc
 from colossalai.context.parallel_mode import ParallelMode
 
+
 class ColoLinear(ColoModule):
+
     def __init__(self):
         super(ColoLinear, self).__init__()
         self._register_shard_params(['weight', 'bias'])
-        
+
     def register(self, compute_pattern):
         if not compute_pattern in self._allowed_patterns:
             if ComputePattern.TP1D == compute_pattern:
                 self._set_TP1D()
-        
+
     def _set_TP1D(self):
         # TP1D Row Linear
         _compute_pattern = ComputePattern.TP1D
         self._register_allowed_patterns(
             compute_pattern=_compute_pattern,
             dist_specs={
-                'weight': distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-                'bias': None
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [-1],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'bias':
+                    None
             },
             mode='row',
         )
@@ -29,8 +34,12 @@ class ColoLinear(ColoModule):
         self._register_allowed_patterns(
             compute_pattern=_compute_pattern,
             dist_specs={
-                'weight': distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
-                'bias': distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0], [gpc.get_world_size(ParallelMode.PARALLEL_1D)])
+                'weight':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)]),
+                'bias':
+                    distspec.shard(gpc.get_group(ParallelMode.PARALLEL_1D), [0],
+                                   [gpc.get_world_size(ParallelMode.PARALLEL_1D)])
             },
             mode='col',
         )
diff --git a/colossalai/nn/optimizer/__init__.py b/colossalai/nn/optimizer/__init__.py
index 14cb01c24..f9a2bc98f 100644
--- a/colossalai/nn/optimizer/__init__.py
+++ b/colossalai/nn/optimizer/__init__.py
@@ -7,7 +7,9 @@ from .lamb import Lamb
 from .lars import Lars
 from .cpu_adam import CPUAdam
 from .hybrid_adam import HybridAdam
+from .colo_optimizer import ColoOptimizer
 
 __all__ = [
-    'ColossalaiOptimizer', 'FusedLAMB', 'FusedAdam', 'FusedSGD', 'Lamb', 'Lars', 'CPUAdam', 'HybridAdam', 'CPU_ADAM_CNT'
+    'ColossalaiOptimizer', 'FusedLAMB', 'FusedAdam', 'FusedSGD', 'Lamb', 'Lars', 'CPUAdam', 'HybridAdam',
+    'CPU_ADAM_CNT', 'ColoOptimizer'
 ]
diff --git a/colossalai/tensor/optim/colo_optimizer.py b/colossalai/nn/optimizer/colo_optimizer.py
similarity index 100%
rename from colossalai/tensor/optim/colo_optimizer.py
rename to colossalai/nn/optimizer/colo_optimizer.py
diff --git a/colossalai/tensor/__init__.py b/colossalai/tensor/__init__.py
index 008183280..a13cfbec1 100644
--- a/colossalai/tensor/__init__.py
+++ b/colossalai/tensor/__init__.py
@@ -1,21 +1,14 @@
 from .spec import ComputePattern, ParallelAction, TensorSpec
-from .op_wrapper import (
-    colo_op_impl,)
+
 from .colo_tensor import ColoTensor
 from .colo_parameter import ColoParameter
 from .utils import convert_parameter, named_params_with_colotensor
-from ._ops import *
-from .optim.colo_optimizer import ColoOptimizer
 from . import distspec
 from .dist_spec_mgr import DistSpecManager
 from .param_op_hook import ParamOpHook, use_param_op_hooks
 from .chunk import ChunkManager, TensorState
-from .module_utils import register_colo_module, is_colo_module, get_colo_module, init_colo_module, check_colo_module
-from .modules import ColoLinear, ColoEmbedding
 
 __all__ = [
-    'ColoTensor', 'convert_parameter', 'colo_op_impl', 'ComputePattern', 'TensorSpec', 'ParallelAction',
-    'named_params_with_colotensor', 'ColoOptimizer', 'ColoParameter', 'distspec', 'DistSpecManager',
-    'register_colo_module', 'is_colo_module', 'get_colo_module', 'init_colo_module', 'check_colo_module', 'ColoLinear',
-    'ColoEmbedding', 'ParamOpHook', 'use_param_op_hooks', 'ChunkManager', 'TensorState'
+    'ColoTensor', 'convert_parameter', 'ComputePattern', 'TensorSpec', 'ParallelAction', 'named_params_with_colotensor',
+    'ColoParameter', 'distspec', 'DistSpecManager', 'ParamOpHook', 'use_param_op_hooks', 'ChunkManager', 'TensorState'
 ]
diff --git a/colossalai/tensor/colo_parameter.py b/colossalai/tensor/colo_parameter.py
index 74eaa0b97..ea040aca3 100644
--- a/colossalai/tensor/colo_parameter.py
+++ b/colossalai/tensor/colo_parameter.py
@@ -1,9 +1,9 @@
-from .colo_tensor import ColoTensor
-from .const import TensorType
+from colossalai.tensor.colo_tensor import ColoTensor
+from colossalai.tensor.const import TensorType
 import torch
 from colossalai.tensor import TensorSpec, distspec
 from copy import copy
-from .param_op_hook import _ParamOpHookWrapper, PreFwdPostBwd, PostFwdPreBwd
+from colossalai.tensor.param_op_hook import _ParamOpHookWrapper, PreFwdPostBwd, PostFwdPreBwd
 from typing import Optional
 
 
diff --git a/colossalai/tensor/dist_spec_mgr.py b/colossalai/tensor/dist_spec_mgr.py
index a3e229899..a1722ff8f 100644
--- a/colossalai/tensor/dist_spec_mgr.py
+++ b/colossalai/tensor/dist_spec_mgr.py
@@ -1,11 +1,29 @@
 from colossalai.tensor.distspec import _DistSpec
-from colossalai.nn.layer.utils import divide
+# from colossalai.nn.layer.utils import divide
 from numpy import prod
 from contextlib import contextmanager
 import torch
 import torch.distributed as dist
 
 
+# TODO(jiaruifang) circle import, move the divide to colossalai.commons.
+# colossalai.tensor shall not import any submodule from colossal.nn
+def divide(numerator, denominator):
+    """Only allow exact division.
+
+    Args:
+        numerator (int): Numerator of the division.
+        denominator (int): Denominator of the division.
+
+    Returns:
+        int: the result of exact division.
+    """
+    assert denominator != 0, 'denominator can not be zero'
+    assert numerator % denominator == 0, \
+        '{} is not divisible by {}'.format(numerator, denominator)
+    return numerator // denominator
+
+
 class TransformDistSpec(torch.autograd.Function):
 
     @staticmethod
diff --git a/colossalai/tensor/optim/__init__.py b/colossalai/tensor/optim/__init__.py
deleted file mode 100644
index e69de29bb..000000000
diff --git a/colossalai/tensor/utils.py b/colossalai/tensor/utils.py
index 5abce8ca1..91cc3aeb9 100644
--- a/colossalai/tensor/utils.py
+++ b/colossalai/tensor/utils.py
@@ -1,10 +1,8 @@
 import torch
 
-from colossalai.tensor.colo_tensor import ColoTensor
-
 from typing import Iterator, Tuple, Union
 import torch.nn as nn
-from colossalai.tensor import ColoTensor
+from colossalai.tensor.colo_tensor import ColoTensor
 
 
 # The function is credited to PyTorch Team
diff --git a/colossalai/utils/model/colo_init_context.py b/colossalai/utils/model/colo_init_context.py
index 807f9034a..87489f0ce 100644
--- a/colossalai/utils/model/colo_init_context.py
+++ b/colossalai/utils/model/colo_init_context.py
@@ -1,11 +1,12 @@
 from .utils import InsertPostInitMethodToModuleSubClasses
 import torch
-from colossalai.tensor import ColoTensor, ColoParameter, register_colo_module, init_colo_module, \
+from colossalai.tensor import ColoTensor, ColoParameter
+
+from colossalai.nn import register_colo_module, init_colo_module, \
     ColoLinear, ColoEmbedding
-import types
 
 from torch import nn
-from typing import Iterator, Tuple, Union, Optional
+from typing import Iterator, Tuple, Union
 
 # find named_params includes replica
 
@@ -24,6 +25,7 @@ def _named_params_with_replica(
             name = mod_prefix + ('.' if mod_prefix else '') + name
             yield name, val
 
+
 def ColoModulize(module):
     """
     Replacing the parameters() and named_parameters() with our customized ones
diff --git a/tests/test_tensor/test_hybrid_device.py b/tests/test_tensor/test_hybrid_device.py
index 4a7a596a8..f1dcd8f20 100644
--- a/tests/test_tensor/test_hybrid_device.py
+++ b/tests/test_tensor/test_hybrid_device.py
@@ -1,9 +1,12 @@
 from colossalai.utils import free_port, ColoInitContext, get_current_device
 from colossalai.testing import rerun_if_address_is_in_use
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, init_colo_module
+from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction
+
 from functools import partial
 from colossalai.core import global_context as gpc
 from colossalai.context import ParallelMode
+
+from colossalai.nn import init_colo_module
 from colossalai.nn.parallel import ColoDDP
 
 import colossalai
@@ -11,12 +14,14 @@ import torch
 import torch.multiprocessing as mp
 import pytest
 
+
 class Net(torch.nn.Module):
+
     def __init__(self):
         super(Net, self).__init__()
         self.embed = torch.nn.Embedding(20, 4)
         self.proj = torch.nn.Linear(4, 8)
-    
+
     def forward(self, x):
         # move input to cpu and restore output
         current_dev = x.device
@@ -27,6 +32,7 @@ class Net(torch.nn.Module):
         x = self.proj(x)
         return x
 
+
 def run_hybrid_device(use_ddp):
     with ColoInitContext(device=get_current_device()):
         model = Net()
@@ -36,7 +42,6 @@ def run_hybrid_device(use_ddp):
         model = ColoDDP(model)
         real_model = model.module
 
-
     print(f'embedding weight size: {real_model.embed.weight.size()} | device: {real_model.embed.weight.device}')
     #print(f'linear weight size: {real_model.proj.weight.size()} | device: {real_model.proj.weight.device}')
     parallel_action = ParallelAction(ComputePattern.TP1D)
@@ -49,11 +54,12 @@ def run_hybrid_device(use_ddp):
 
     print(f'embedding weight size: {real_model.embed.weight.size()} | new device: {real_model.embed.weight.device}')
     #print(f'linear weight size: {real_model.proj.weight.size()} | new device: {real_model.proj.weight.device}')
-    
+
     data = torch.randint(low=0, high=20, size=(16,), device=get_current_device())
     out = model(data)
     out.sum().backward()
 
+
 def run_dist(rank, world_size, port, use_ddp):
     if use_ddp and world_size == 1:
         return
@@ -62,6 +68,7 @@ def run_dist(rank, world_size, port, use_ddp):
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_hybrid_device(use_ddp)
 
+
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 4])
 @pytest.mark.parametrize('use_ddp', [False, True])
@@ -71,5 +78,6 @@ def _test_hybrid_device(world_size, use_ddp):
     run_func = partial(run_dist, world_size=world_size, port=free_port(), use_ddp=use_ddp)
     mp.spawn(run_func, nprocs=world_size)
 
+
 if __name__ == '__main__':
-    _test_hybrid_device(1, False)
\ No newline at end of file
+    _test_hybrid_device(1, False)
diff --git a/tests/test_tensor/test_model.py b/tests/test_tensor/test_model.py
index 682146e1a..019152aa6 100644
--- a/tests/test_tensor/test_model.py
+++ b/tests/test_tensor/test_model.py
@@ -10,9 +10,10 @@ from colossalai.utils.cuda import get_current_device
 from colossalai.utils import free_port
 from colossalai.utils import ColoInitContext
 from colossalai.tensor import distspec, named_params_with_colotensor, TensorSpec, ComputePattern, \
-    ParallelAction, ColoTensor, ColoOptimizer, DistSpecManager
+    ParallelAction, ColoTensor, DistSpecManager
 from colossalai.context import ParallelMode
 from colossalai.core import global_context as gpc
+from colossalai.nn.optimizer import ColoOptimizer
 from functools import partial
 from _utils import set_seed
 
diff --git a/tests/test_tensor/test_module_spec.py b/tests/test_tensor/test_module_spec.py
index 771662571..0a4e7c9d2 100644
--- a/tests/test_tensor/test_module_spec.py
+++ b/tests/test_tensor/test_module_spec.py
@@ -1,24 +1,28 @@
 from copy import copy
-from colossalai.utils.cuda import get_current_device
-from colossalai.utils.model.colo_init_context import ColoInitContext
-import torch
-from colossalai.context.parallel_mode import ParallelMode
-from colossalai.tensor import ColoTensor, distspec
-
+import pytest
 from functools import partial
 
-import colossalai
-import pytest
 import torch
 import torch.multiprocessing as mp
-import torch.nn.functional as F
+
+from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction
+from colossalai.nn import init_colo_module, check_colo_module
+from _utils import tensor_equal, tensor_shard_equal, set_seed
+
+import colossalai
+from colossalai.utils.cuda import get_current_device
+from colossalai.utils.model.colo_init_context import ColoInitContext
+
+from colossalai.context.parallel_mode import ParallelMode
+from colossalai.tensor import distspec
+
 from colossalai.testing import rerun_if_address_is_in_use
 from colossalai.utils import free_port
 from colossalai.core import global_context as gpc
-from colossalai.tensor import TensorSpec, ComputePattern, ParallelAction, DistSpecManager, register_colo_module, init_colo_module, check_colo_module
-from _utils import tensor_equal, tensor_shard_equal, set_seed
+
 from tests.components_to_test.registry import non_distributed_component_funcs
 
+
 def run_model_with_spec(mode, model_name):
     get_components_func = non_distributed_component_funcs.get_callable(model_name)
     model_builder, train_dataloader, test_dataloader, optimizer_class, criterion = get_components_func()
@@ -27,7 +31,7 @@ def run_model_with_spec(mode, model_name):
     set_seed(1)
     with ColoInitContext(device=get_current_device()):
         model = model_builder(checkpoint=False)
-    
+
     if rank == 0:
         model_seq = model_builder(checkpoint=False)
         model_seq = model_seq.cuda()
@@ -103,15 +107,16 @@ def run_model_with_spec(mode, model_name):
         if i > 3:
             break
 
+
 def run_linear_with_spec(mode):
     with ColoInitContext(device=get_current_device()):
         model = torch.nn.Linear(4, 8)
 
     model_handy = copy(model)
-    
+
     parallel_action = ParallelAction(ComputePattern.TP1D)
     init_colo_module(model, parallel_action, recursive=True, mode=mode)
-    
+
     x = torch.rand(2, 4).cuda()
     out = model(x)
     colo_out = model_handy(x)
@@ -122,6 +127,7 @@ def run_linear_with_spec(mode):
     assert tensor_shard_equal(model.weight.grad, model_handy.weight.grad)
     assert tensor_shard_equal(model.bias.grad, model_handy.bias.grad)
 
+
 def run_check_shared_param():
     from transformers import BertForMaskedLM, BertConfig
     hidden_dim = 8
@@ -157,12 +163,14 @@ def run_check_shared_param():
     except Exception as e:
         assert 'incorrectly sharded' in str(e)
 
+
 def run_dist(rank, world_size, port):
     config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_linear_with_spec('col')
     run_linear_with_spec('row')
 
+
 def run_dist_model(rank, world_size, port):
     config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
@@ -170,11 +178,13 @@ def run_dist_model(rank, world_size, port):
         run_model_with_spec('col', model_name)
         run_model_with_spec('row', model_name)
 
+
 def run_dist_check(rank, world_size, port):
     config = dict(parallel=dict(tensor=dict(mode="1d", size=world_size),))
     colossalai.launch(config=config, rank=rank, world_size=world_size, host='localhost', port=port, backend='nccl')
     run_check_shared_param()
 
+
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 4])
 @rerun_if_address_is_in_use()
@@ -182,6 +192,7 @@ def test_module_linear_1d(world_size):
     run_func = partial(run_dist, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)
 
+
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 4])
 @rerun_if_address_is_in_use()
@@ -189,6 +200,7 @@ def test_module_model(world_size):
     run_func = partial(run_dist_model, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)
 
+
 @pytest.mark.dist
 @pytest.mark.parametrize('world_size', [1, 2])
 @rerun_if_address_is_in_use()
@@ -196,5 +208,6 @@ def test_module_check(world_size):
     run_func = partial(run_dist_check, world_size=world_size, port=free_port())
     mp.spawn(run_func, nprocs=world_size)
 
+
 if __name__ == '__main__':
-    test_module_check(2)
\ No newline at end of file
+    test_module_check(2)