diff options
Diffstat (limited to 'ep_run')
| -rw-r--r-- | ep_run/casc_eq_train.py | 1 |
1 files changed, 1 insertions, 0 deletions
diff --git a/ep_run/casc_eq_train.py b/ep_run/casc_eq_train.py index 60b6113..0cd5521 100644 --- a/ep_run/casc_eq_train.py +++ b/ep_run/casc_eq_train.py @@ -908,6 +908,7 @@ for step in range(start_step, args.steps + 1): with torch.no_grad(): GOV['w_rms'] = float(sum(p.float().pow(2).mean().sqrt() for p in all_params) / len(all_params)) _aux = {'gn': GOV.get('gn'), 'drift': GOV.get('drift'), 'sig': GOV.get('sig'), + 'res': GOV.get('res'), 'rho': GOV.get('rho'), 'clip_norm': GOV.get('clip_norm'), 'clip_fired': (None if GOV.get('clip_norm') is None else float(GOV['clip_norm'] > 1.0))} for k in ('act_rms', 'w_rms', 'mag_ep', 'mag_bp'): if GOV.get(k) is not None: _aux[k] = GOV[k] |