from colossalai.tensor.d_tensor.api import to_distributed_tensor