import pytest
import torch

from colossalai.utils.cuda import get_current_device
from colossalai.utils.model.colo_init_context import ColoInitContext