torch.utils.module_tracker#

Created On: May 04, 2024 | Last Updated On: Jun 11, 2025

This utility can be used to track the current position inside antorch.nn.Module hierarchy.It can be used within other tracking tools to be able to easily associate measured quantities to user-friendly names. This is used in particular in the FlopCounterMode today.

classtorch.utils.module_tracker.ModuleTracker[source]#

ModuleTracker is a context manager that tracks the nn.Module hierarchy during executionso that other system can query which Module is currently being executed (or its backward is beingexecuted).

You can access theparents attribute on this context manager to get the set of all theModules currently being executed via their fqn (fully qualified name, also used as the key withinthe state_dict).You can access theis_bw attribute to know if you are currently running in backward or not.

Note thatparents is never empty and always contains the “Global” key. Theis_bw flagwill remainTrue after the forward until another Module is executed. If you need it to bemore accurate, please submit an issue requesting this. Adding a map from fqn to the module instanceis possible but not done yet, please submit an issue requesting this if you need it.

Example usage

mod=torch.nn.Linear(2,2)withModuleTracker()astracker:# Access anything during the forward passdefmy_linear(m1,m2,bias):print(f"Current modules:{tracker.parents}")returntorch.mm(m1,m2.t())+biastorch.nn.functional.linear=my_linearmod(torch.rand(2,2))