github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-31 08:34:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
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
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 time import time
from typing import List, Optional
from typing import List, Optional, Type
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.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):