ColossalAI/tests/test_tensor/test_context.py

import pytest
import torch

from colossalai.utils.cuda import get_current_device
from colossalai.utils.model.colo_init_context import ColoInitContext
[tensor] refactor colo-tensor (#992) * refactor colo-tensor and update linear op * polish code * polish code * update ops and unit tests * update unit tests * polish code * rename dist_spec module * polish code * polish code * remove unneeded import * fix pipelinable 2022-05-19 04:44:59 +00:00			`import pytest`
Init Conext supports lazy allocate model memory (#842) 2022-04-22 10:03:35 +00:00			`import torch`

colo init context add device attr. (#866) 2022-04-25 06:24:26 +00:00			`from colossalai.utils.cuda import get_current_device`
[utils] remove lazy_memory_allocate from ColoInitContext (#1844) 2022-11-09 03:50:33 +00:00			`from colossalai.utils.model.colo_init_context import ColoInitContext`