(self, module: torch.nn.Module, name: str)
| 133 | return module |
| 134 | |
| 135 | def _set_context(self, module: torch.nn.Module, name: str) -> None: |
| 136 | # Iterate over all attributes of the hook to see if any of them have the type `StateManager`. If so, call `set_context` on them. |
| 137 | for attr_name in dir(self): |
| 138 | attr = getattr(self, attr_name) |
| 139 | if isinstance(attr, StateManager): |
| 140 | attr.set_context(name) |
| 141 | return module |
| 142 | |
| 143 | |
| 144 | class HookFunctionReference: |
no test coverage detected