log_memory_stats

torchtune.utils.log_memory_stats(stats: Dict[str, float]) → 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.utils.get_memory_stats().

Parameters:

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