log_memory_stats¶
- torchtune.training.log_memory_stats(stats: Dict[str, float], message: str = 'Memory stats after model init:') None [source]¶
Logs a dict containing memory stats to the logger.
stats
should contain the fieldspeak_memory_active
,peak_memory_alloc
, andpeak_memory_reserved
as returned bytorchtune.training.get_memory_stats()
.