/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/booster/plugin/pp_plugin_base.py
22 строки
559 B
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
from abc import abstractmethod from typing import Any, Callable, Iterator, Optional import torch from colossalai.interface import ModelWrapper, OptimizerWrapper from .plugin_base import Plugin class PipelinePluginBase(Plugin): @abstractmethod def execute_pipeline( self, data_iter: Iterator, model: ModelWrapper, criterion: Callable[[Any, Any], torch.Tensor], optimizer: Optional[OptimizerWrapper] = None, return_loss: bool = True, return_outputs: bool = False, ) -> dict: pass