Commit abe4f971 authored by 梁爽's avatar 梁爽 Committed by binmakeswell
Browse files

[NFC] polish colossalai/booster/plugin/low_level_zero_plugin.py code style (#4256)


Co-authored-by: default avatarsupercooledith <893754954@qq.com>
parent c614a99d
...@@ -208,10 +208,7 @@ class LowLevelZeroPlugin(DPPluginBase): ...@@ -208,10 +208,7 @@ class LowLevelZeroPlugin(DPPluginBase):
if optimizer is not None and \ if optimizer is not None and \
not isinstance(optimizer, OptimizerWrapper): not isinstance(optimizer, OptimizerWrapper):
optimizer = LowLevelZeroOptimizer(model.unwrap(), optimizer = LowLevelZeroOptimizer(model.unwrap(), optimizer, self.zero_optim_config, self.optim_kwargs,
optimizer,
self.zero_optim_config,
self.optim_kwargs,
self.verbose) self.verbose)
return model, optimizer, criterion, dataloader, lr_scheduler return model, optimizer, criterion, dataloader, lr_scheduler
......
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