/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
tests/kit/model_zoo/custom/base.py
26 строк
832 B
Hongxin Liu
[test] merge old components to test to model zoo (#4945)
20 окт 2023, 05:35
Не верифицирован
20 окт 2023, 05:35
b8e770c
Код
Авторство
О чём код?
import torch.nn as nn from torch.utils.checkpoint import checkpoint class CheckpointModule(nn.Module): def __init__(self, checkpoint: bool = False): super().__init__() self.checkpoint = checkpoint self._use_checkpoint = checkpoint def _forward(self, *args, **kwargs): raise NotImplementedError("CheckpointModule should implement _forward method instead of origin forward") def forward(self, *args, **kwargs): if self._use_checkpoint: return checkpoint(self._forward, *args, **kwargs) else: return self._forward(*args, **kwargs) def train(self, mode: bool = True): self._use_checkpoint = self.checkpoint return super().train(mode=mode) def eval(self): self._use_checkpoint = False return super().eval()