Skip to content

GitLab

Commit c4a226b7 authored by Camille Zhong's avatar Camille Zhong Committed by binmakeswell
Browse files

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

parent 00778abc
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 7 deletions
colossalai/gemini/tensor_placement_policy.py
View file @ c4a226b7
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.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 colossalai.gemini.memory_tracer import MemStatsCollector
from typing import Type from colossalai.gemini.stateful_tensor import StatefulTensor
import functools 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
class TensorPlacementPolicy(ABC): class TensorPlacementPolicy(ABC):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment