Commit b366f1d9 authored by Jianghai's avatar Jianghai Committed by binmakeswell
Browse files

[NFC] Fix format for mixed precision (#4253)

* [NFC] polish colossalai/booster/mixed_precision/mixed_precision_base.py code style
parent 02192a63
...@@ -13,10 +13,11 @@ class MixedPrecision(ABC): ...@@ -13,10 +13,11 @@ class MixedPrecision(ABC):
""" """
@abstractmethod @abstractmethod
def configure(self, def configure(
model: nn.Module, self,
optimizer: Optional[Optimizer] = None, model: nn.Module,
criterion: Optional[Callable] = None, optimizer: Optional[Optimizer] = None,
) -> Tuple[nn.Module, OptimizerWrapper, Callable]: criterion: Optional[Callable] = None,
) -> Tuple[nn.Module, OptimizerWrapper, Callable]:
# TODO: implement this method # TODO: implement this method
pass pass
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment