Add param counting to train script (#49)

EleutherAI · Oct 16, 2023 · 3c7bdf5 · 3c7bdf5
1 parent 22e0657
commit 3c7bdf5
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/aria/train.py b/aria/train.py
@@ -244,6 +244,11 @@ def _bench():
  break
 
  flop_counter = FlopCounterMode(display=False)
+ logger.info(
+ f"Model has "
+ f"{'{:,}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad))} "
+ "parameters"
+ )
  logger.info("Profiling FLOP/s")
  _bench()