Spaces:
Running
Running
feat: log num_params
Browse files- 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(
|