Fixup Zero3 + `save_model` (#3146)

Commit

1 year ago

Fixup Zero3 + `save_model` (#3146) * Fixup + test * Easier diff * Move os.makedirs to under return statement