torchtnt.utils.module_summary.prune_module_summary

torchtnt.utils.module_summary.prune_module_summary(module_summary: ModuleSummary, *, max_depth: int) → None

Prune the module summaries that are deeper than max_depth in the module summary tree. The ModuleSummary object is prunned inplace.

Parameters:

• module_summary – Root module summary to prune.
• max_depth – The maximum depth of module summaries to keep.

Raises:

ValueError – If max_depth is an int less than 1