githubmirror/ColossalAI

colossalai/interface/pretrained.py

16 строк333 B

Hongxin Liu

[lazy] support from_pretrained (#4801)

26 сен 2023, 06:04

26 сен 2023, 06:044965c0d

from typing import Optional

from torch.nn import Module

__all__ = [
    "get_pretrained_path",
    "set_pretrained_path",
]


def get_pretrained_path(model: Module) -> Optional[str]:
    return getattr(model, "_pretrained", None)


def set_pretrained_path(model: Module, path: str) -> None:
    setattr(model, "_pretrained", path)