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 time import time
from typing import List, Optional
from typing import List, Optional, Type
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 typing import Type
import functools
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
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