Browse Source

Update low_level_zero_plugin.py

pull/6029/head
flybird11111 3 months ago committed by GitHub
parent
commit
9e767643dd
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 2
      colossalai/booster/plugin/low_level_zero_plugin.py

2
colossalai/booster/plugin/low_level_zero_plugin.py

@ -349,7 +349,7 @@ class LowLevelZeroPlugin(DPPluginBase):
verbose: bool = False, verbose: bool = False,
cast_inputs: bool = True, cast_inputs: bool = True,
fp8_communication: bool = False, fp8_communication: bool = False,
use_fp8: bool = False, use_fp8: bool = False
) -> None: ) -> None:
super().__init__() super().__init__()
assert stage in (1, 2), f"LowLevelZeroPlugin only supports stage 1/2 training" assert stage in (1, 2), f"LowLevelZeroPlugin only supports stage 1/2 training"

Loading…
Cancel
Save