Log model

This commit is contained in:
Srini Iyer 2025-02-06 00:26:37 +00:00
parent 7cf8fab49b
commit 162b99b4a3

View file

@ -325,6 +325,7 @@ def train(args: TrainArgs):
# log model size # log model size
logger.info(model)
logger.info(f"Model size: {model_param_count:,} total parameters") logger.info(f"Model size: {model_param_count:,} total parameters")
gpu_memory_monitor = GPUMemoryMonitor("cuda") gpu_memory_monitor = GPUMemoryMonitor("cuda")