log_memory_stats

torchtune.training.log_memory_stats(stats: Dict[str, float], message: str = 'Memory stats after model init:') → None

Logs a dict containing memory stats to the logger. stats should contain the fields peak_memory_active, peak_memory_alloc, and peak_memory_reserved as returned by torchtune.training.get_memory_stats().

Parameters:

• stats (Dict[str, float]) – A dictionary containing the peak memory active, peak memory allocated, and peak memory reserved stats.

• message (str) – An optional message to prepend to the log output. Defaults to “Memory stats after model init:”