Repository: pytorch
Issue: Can we apply Weight normalization hook to a method other than `forward`?
Labels: module: nn,triaged,enhancement,module: norms and normalization

## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

Weight normalization, as implemented in pytorch, is a forward hook for `nn.Module`. It is called before every call to the `forward` method of the module. But I think it's common for a module to have more computation methods other than `forward`. But say, that I have a method that directly accesses the `weight` and `bias` without invoking the forward method, for example, to implement a special method for a Conv1d Module, which enables it to perform `step-in-step-out` behavior.

Conv1d: 
    forward: x:(batch_size, in_channels, time_steps) ->  x:(batch_size, in_channels, time_steps) with causal padding. i.e. padding with (receptive_field - 1, 0)
    step_in_step_out: x(batch_size, in_channels, 1) -> x:(batch_size, in_channels, 1), by caching a buffer with shape (batch_size, in_channels, receptive_field - 1) and update the buffer at every call to this method. We can implement this by calling `F.conv1d(step_input, self.weight, self.bias, ) with padding=0.

But if the Conv1d has a weight_norm hook. Then calling `step_in_step_out` would not invoke the hook. If we load from saved state_dict( only weight_v, weight_g and bias in state dict), and immediately call `step_in_step_out`, the `weight` is not updated so the old weight is used, which is unexpected.

Are there better ways to handle this problem? Can I explicitly call the forward hook for another method other than `forward`?
Thank you.


## To Reproduce

Steps to reproduce the behavior:

```python
import torch
from torch import nn
from torch.nn import functional as F 

class MyLinear(nn.Linear):
    def special_forward(self, x):
        return F.linear(x, self.weight, self.bias)

x = torch.randn([2, 3])
linear_layer = nn.utils.weight_norm(MyLinear(3, 4))
y = linear_layer(x)
print(y)

linear_layer2 = nn.utils.weight_norm(MyLinear(3, 4))
linear_layer2.load_state_dict(linear_layer.state_dict())
y2 = linear_layer2.special_forward(x)
print(y2)

print("y1 != y2")
```

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

## Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

 - PyTorch Version (e.g., 1.0): 1.7
 - OS (e.g., Linux):
 - How you installed PyTorch (`conda`, `pip`, source):
 - Build command you used (if compiling from source):
 - Python version:
 - CUDA/cuDNN version:
 - GPU models and configuration:
 - Any other relevant information:

## Additional context

<!-- Add any other context about the problem here. -->


cc @albanD @mruberry