mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-05-03 05:58:09 +00:00
079bf3cb26
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
23 lines
559 B
Python
23 lines
559 B
Python
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
|