3.7.4. unit_scaling.transforms.utils.torch_nn_modules_to_user_modules

unit_scaling.transforms.utils.torch_nn_modules_to_user_modules(mod: Module) → None

Convert torch.nn.module classes to trivial_subclass versions.

By default TorchDynamo doesn’t recurse into torch.nn modules or torch.nn.functional functions when capturing the FX graph.

This function makes torch.nn modules into user modules.

To use this with a torch.nn.Module the typical use case is to call module = torch_nn_modules_to_user_modules(module).