Commit d68ec712 authored by Michael Carilli's avatar Michael Carilli
Browse files

test cleanup

parent 848c777d
...@@ -758,13 +758,8 @@ class TestMultipleModelsOptimizersLosses(unittest.TestCase): ...@@ -758,13 +758,8 @@ class TestMultipleModelsOptimizersLosses(unittest.TestCase):
inj_model.weight1.grad[0] = float('inf') inj_model.weight1.grad[0] = float('inf')
if i != inject_inf: if i != inject_inf:
if opt_level == "O2" and not materialize_master_grads: master_params = list(amp.master_params(optimizer0)) + \
master_params = list(model0.parameters()) + \ list(amp.master_params(optimizer1))
list(model1.parameters()) + \
list(model2.parameters())
else:
master_params = list(amp.master_params(optimizer0)) + \
list(amp.master_params(optimizer1))
for param, reference_grad in zip(master_params, for param, reference_grad in zip(master_params,
reference_grads[what_got_skipped(inject_inf, reference_grads[what_got_skipped(inject_inf,
which_backward, which_model)][unskipped]): which_backward, which_model)][unskipped]):
......
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