mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-07-30 15:00:39 +00:00
parent
1ee947f617
commit
2d40759a53
@ -15,7 +15,7 @@ from colossalai.nn import LinearWarmupLR
|
||||
from colossalai.trainer import Trainer, hooks
|
||||
from colossalai.utils import colo_set_process_memory_fraction, is_using_pp
|
||||
from colossalai.utils.timer import MultiTimer
|
||||
from colossalai.zero.init_ctx import ZeroInitContext
|
||||
from colossalai.zero.legacy.init_ctx import ZeroInitContext
|
||||
|
||||
|
||||
def calc_local_model_size(model: torch.nn.Module):
|
||||
|
Loading…
Reference in New Issue
Block a user