github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish tensor_placement_policy.py code style (#3265)

pull/3313/head
Camille Zhong 2023-03-27 22:03:46 +08:00 committed by binmakeswell
parent 00778abc48
commit c4a226b729
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
colossalai/gemini/tensor_placement_policy.py
View File

@ -1,16 +1,16 @@
import functools
from abc import ABC, abstractmethod from abc import ABC, abstractmethod
from time import time from time import time
from typing import List, Optional from typing import List, Optional, Type
import torch import torch
from colossalai.gemini.memory_tracer import MemStatsCollector
from colossalai.gemini.stateful_tensor import StatefulTensor
from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage
from colossalai.utils import get_current_device from colossalai.utils import get_current_device
from colossalai.utils.memory import colo_device_memory_capacity from colossalai.utils.memory import colo_device_memory_capacity
from colossalai.gemini.tensor_utils import colo_model_data_tensor_move_inline, colo_tensor_mem_usage
from colossalai.gemini.stateful_tensor import StatefulTensor
from colossalai.gemini.memory_tracer import MemStatsCollector
from typing import Type
import functools
class TensorPlacementPolicy(ABC): class TensorPlacementPolicy(ABC):