Repository: pytorch
Issue: torch.nn.utils.prune.remove reorders the parameters of a module unexpectedly
Labels: triaged,module: pruning

### 🐛 Describe the bug

PyTorch has a pruning tool called `torch.nn.utils.prune`. It includes a function `torch.nn.utils.prune.remove`, which should restore the model parameter structure back to normal, except it keeps the effect of the pruning.

Since nn.module.state_dict() returns an **Ordered**Dict, I would expect that there is some internal ordering of the parameters of a module, and this ordering should NOT be changed by pruning (as long as I called `torch.nn.utils.prune.remove` to remove the  reparameterization).

However, during some testing, I found that pruning and `torch.nn.utils.prune.remove` **does** change the order of the parameters. 

This means, pruning would break some code that depends on this ordering. (I'm trying to prune [this model](https://github.com/openai/guided-diffusion), which depends on the ordering to save the checkpoint during training.)

For example, in the following code, the `weight` and `bias` of this Conv2d module is swapped unexpectedly:

```python
import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

module=nn.Conv2d(1, 6, 3)
print(list(module.state_dict().keys())) #output: ['weight', 'bias']

prune.random_unstructured(module, name="weight", amount=0.3)
prune.remove(module, name="weight")

print(list(module.state_dict().keys())) #output: ['bias', 'weight']
```

Expected output:
```
['weight', 'bias']
['weight', 'bias']
```

Actual output:
```
['weight', 'bias']
['bias', 'weight']
```

### Versions

```
Collecting environment information...
PyTorch version: 1.9.1+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: CentOS Linux 7 (Core) (x86_64)
GCC version: (GCC) 4.8.5 20150623 (Red Hat 4.8.5-44)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.17

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.49.1.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.1
[pip3] torch==1.9.1+cu111
[pip3] torchaudio==0.9.1
[pip3] torchvision==0.10.1+cu111
[conda] _tflow_select             2.3.0                       mkl  
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               10.2.89              hfd86e86_1  
[conda] mkl                       2021.4.0           h06a4308_640  
[conda] mkl-service               2.4.0            py38h7f8727e_0  
[conda] mkl_fft                   1.3.1            py38hd3c417c_0  
[conda] mkl_random                1.2.2            py38h51133e4_0  
[conda] numpy                     1.23.1           py38h6c91a56_0  
[conda] numpy-base                1.23.1           py38ha15fc14_0  
[conda] tensorflow                2.8.2           mkl_py38hb41d75a_0  
[conda] tensorflow-base           2.8.2           mkl_py38hf890080_0  
[conda] torch                     1.9.1+cu111              pypi_0    pypi
[conda] torchaudio                0.9.1                    pypi_0    pypi
[conda] torchvision               0.10.1+cu111             pypi_0    pypi
```