mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-05 02:51:59 +00:00
[lazy] support from_pretrained (#4801)
* [lazy] patch from pretrained * [lazy] fix from pretrained and add tests * [devops] update ci
This commit is contained in:
16
colossalai/interface/pretrained.py
Normal file
16
colossalai/interface/pretrained.py
Normal file
@@ -0,0 +1,16 @@
|
||||
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)
|
Reference in New Issue
Block a user