boris commited on
Commit
1f57ad7
1 Parent(s): 8654dc9

feat: log num_params

Browse files
Files changed (1) hide show
  1. tools/train/train.py +2 -0
tools/train/train.py CHANGED
@@ -617,6 +617,7 @@ def main():
617
  "len_train_dataset": len_train_dataset,
618
  "len_eval_dataset": len_eval_dataset,
619
  "batch_size_per_update": batch_size_per_update,
 
620
  }
621
  )
622
 
@@ -693,6 +694,7 @@ def main():
693
  c.cleanup(wandb.util.from_human_size("10GB"))
694
 
695
  metadata = dict(state_dict)
 
696
  if eval_metrics is not None:
697
  metadata["eval"] = eval_metrics
698
  artifact = wandb.Artifact(
 
617
  "len_train_dataset": len_train_dataset,
618
  "len_eval_dataset": len_eval_dataset,
619
  "batch_size_per_update": batch_size_per_update,
620
+ "num_params": model.num_params,
621
  }
622
  )
623
 
 
694
  c.cleanup(wandb.util.from_human_size("10GB"))
695
 
696
  metadata = dict(state_dict)
697
+ metadata["num_params"] = model.num_params
698
  if eval_metrics is not None:
699
  metadata["eval"] = eval_metrics
700
  artifact = wandb.Artifact(