fix changed signature

mosaicml · Oct 6, 2023 · a776d7c · a776d7c
1 parent f7d03b6
commit a776d7c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/test_lion8b.py b/tests/test_lion8b.py
@@ -460,7 +460,7 @@ def _set_state_dict_type(model: nn.Module):
 
     # load state dict into the new optimizer
     opt_state_dict_slice = FSDP.optim_state_dict_to_load(
-        opt_state_dict, mod_new, opt_new)
+        optim_state_dict=opt_state_dict, model=mod_new, optim=opt_new)
     opt_new.load_state_dict(opt_state_dict_slice)
 
     new_opt_state_dict = FSDP.optim_state_dict(mod_new, opt_new)