torchscan.utils

torchscan.utils.format_info(module_info, wrap_mode='mid')

Print module summary for an expected input tensor shape

Parameters:

• module_info (dict) – dictionary output of crawl_module

• wrap_mode (str, optional) – wrapping mode

torchscan.utils.aggregate_info(info, max_depth)

Aggregate module information to a maximum depth

Parameters:

• info (dict) – dictionary output of crawl_module

• max_depth (int, optional) – depth at which parent node aggregates children information

Returns:

edited dictionary information

Return type:

dict