Skip to content

Commit

Permalink
Fix final_models order
Browse files Browse the repository at this point in the history
  • Loading branch information
KSGulin committed Feb 8, 2023
1 parent 2b4a23a commit 6e3944f
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion train.py
Original file line number Diff line number Diff line change
Expand Up @@ -481,7 +481,7 @@ def _create_dataloaders():
if RANK in {-1, 0}:
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
best_models = [best] if not sparsification_manager else [ w / f"best_{phase}.pt" for phase in sparsification_manager.passed_phases]
final_models = [last] + best_models
final_models = best_models + [last]
for f in final_models:
if f.exists():
strip_optimizer(f) if not sparsification_manager else sparsification_manager.strip_sparsified_optimizer(f) # strip optimizers
Expand Down

0 comments on commit 6e3944f

Please sign in to comment.