hotfix check ckpts before writing achieved model (#3901)

this killed tinybox green run
This commit is contained in:
chenyu
2024-03-23 17:16:38 -04:00
committed by GitHub
parent 4d566f12b1
commit 24d004a89b

View File

@@ -200,6 +200,7 @@ def train_resnet():
# save model if achieved target
if not achieved and total_top_1 >= target:
if not os.path.exists("./ckpts"): os.mkdir("./ckpts")
fn = f"./ckpts/resnet50.safe"
safe_save(get_state_dict(model), fn)
print(f" *** Model saved to {fn} ***")