Shortcuts

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 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:”

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources