github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-08-07 02:54:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* [NFC] polish colossalai/booster/mixed_precision/mixed_precision_base.py code style
This commit is contained in:
Jianghai 2023-07-17 18:01:30 +08:00 committed by binmakeswell
parent 02192a632e
commit b366f1d99f
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
colossalai/booster/mixed_precision/mixed_precision_base.py
View File

@ -13,7 +13,8 @@ class MixedPrecision(ABC):
""" """
@abstractmethod @abstractmethod
def configure(self, def configure(
self,
model: nn.Module, model: nn.Module,
optimizer: Optional[Optimizer] = None, optimizer: Optional[Optimizer] = None,
criterion: Optional[Callable] = None, criterion: Optional[Callable] = None,