mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-15 22:19:38 +00:00
[pipeline] set optimizer to optional in execute_pipeline (#4630)
* set optimizer to optional in execute_pipeline * arrange device and mixed precision in booster init * fix execute_pipeline in booster.py
This commit is contained in:
@@ -1,5 +1,5 @@
|
||||
from abc import abstractmethod
|
||||
from typing import Any, Callable, Iterator
|
||||
from typing import Any, Callable, Iterator, Optional
|
||||
|
||||
import torch
|
||||
|
||||
@@ -15,7 +15,7 @@ class PipelinePluginBase(Plugin):
|
||||
data_iter: Iterator,
|
||||
model: ModelWrapper,
|
||||
criterion: Callable[[Any, Any], torch.Tensor],
|
||||
optimizer: OptimizerWrapper,
|
||||
optimizer: Optional[OptimizerWrapper] = None,
|
||||
return_loss: bool = True,
|
||||
return_outputs: bool = False) -> dict:
|
||||
pass
|
||||
|
Reference in New Issue
Block a user