| 196 | # Parametrizations come with an inbuilt caching system via the context manager |
| 197 | # ``parametrize.cached()`` |
| 198 | class NoisyParametrization(nn.Module): |
| 199 | def forward(self, X): |
| 200 | print("Computing the Parametrization") |
| 201 | return X |
| 202 | |
| 203 | layer = nn.Linear(4, 4) |
| 204 | parametrize.register_parametrization(layer, "weight", NoisyParametrization()) |