diff options
Diffstat (limited to 'ep_run/casc_bp_train.py')
| -rw-r--r-- | ep_run/casc_bp_train.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/ep_run/casc_bp_train.py b/ep_run/casc_bp_train.py index 7bafc16..0515ec1 100644 --- a/ep_run/casc_bp_train.py +++ b/ep_run/casc_bp_train.py @@ -238,7 +238,7 @@ for step in range(start_step, args.steps + 1): if wb is not None: try: wb.log({'train_ce': loss.item(), 'val_ce': val, 'best': best}, step=step) except Exception: pass - if step % args.save_every == 0: + if step % args.save_every == 0 or step == args.steps: torch.save({'tok': tok.state_dict(), 'pos': pos.state_dict(), 'blocks': blocks.state_dict(), 'wout': (W_out.detach().cpu() if args.olmo2 else None), 'lnf': (ln_f.state_dict() if not isinstance(ln_f, nn.Identity) else None), |