Unverified Commit 9ef05ed1 authored by ver217's avatar ver217 Committed by GitHub
Browse files

try import deepspeed when using zero (#130)

parent b7975d2b
......@@ -29,6 +29,7 @@ def convert_to_zero(model: nn.Module,
:return: (model, optimizer)
:rtype: Tuple
"""
import deepspeed
assert level == 2 or level == 3, 'Only ZERO Optimizer Level 2 and 3 are provided'
model = NaiveAMPModel(model, output_to_fp32=False)
......
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