Repository: pytorch
Issue: Torch serialization does not restore tensors properly when custom __reduce__ is defined
Labels: module: serialization,triaged

## 🐛 Bug

If a `nn.Module` define a custom `__reduce__` method to customize serialization, with the `__reduce__` function returning some tensors (like a state dict), then upon deserialization, the tensors will not contain the right values.

## To Reproduce

```python
import pickle

import torch as th


def stable_instantiate(klass, args, kwargs, state):
    self = klass(*args, **kwargs)
    print("loading", state)
    self.load_state_dict(state)
    return self


def stable_pickle(klass):
    old_init = klass.__init__

    def __init__(self, *args, **kwargs):
        self._init_args_kwargs = (args, kwargs)
        old_init(self, *args, **kwargs)

    def __reduce__(self):
        print("reducing", self.state_dict())
        return (stable_instantiate,
                tuple([klass] + list(self._init_args_kwargs) + [self.state_dict()]))

    klass.__init__ = __init__
    klass.__reduce__ = __reduce__
    return klass


@stable_pickle
class MyModule(th.nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = th.nn.Linear(dim, 1, bias=False)

    def forward(self, x):
        return self.linear(x)


dim = 4
model = MyModule(dim)
model.linear.weight.data[0] = th.arange(dim).float()

print("model weights", model.linear.weight.data)

print("Trying to serialize with torch.save/load")
th.save(model, "tmp.th")
model2 = th.load("tmp.th")
print("loaded model weights", model2.linear.weight.data)

print("Trying to serialize with pickle")
pickle.dump(model, open("tmp.pckl", "wb"))
model2 = pickle.load(open("tmp.pckl", "rb"))
print("unpickled model weights", model2.linear.weight.data)
```

Steps to reproduce the behavior:

1. Save the above code as `bug.py`
1. Run `python3 bug.py`
1. Observe how when calling `torch.load`, the value of the state dict is just random and not what was set before saving.

```bash
model weights tensor([[0., 1., 2., 3.]])
Trying to serialize with torch.save/load
reducing OrderedDict([('linear.weight', tensor([[0., 1., 2., 3.]]))])
loading OrderedDict([('linear.weight', tensor([[-5.6200e-26,  3.0820e-41, -5.5415e-26,  3.0820e-41]]))])
loaded model weights tensor([[-5.6200e-26,  3.0820e-41, -5.5415e-26,  3.0820e-41]])
Trying to serialize with pickle
reducing OrderedDict([('linear.weight', tensor([[0., 1., 2., 3.]]))])
loading OrderedDict([('linear.weight', tensor([[0., 1., 2., 3.]]))])
unpickled model weights tensor([[0., 1., 2., 3.]])
```

## Expected behavior

Expected `torch.save` and `torch.load` to serialize the state dict returned by `__reduce__`. I'm not sure if the problem is at saving or loading but at some point, all the information is lost and a uninitialized tensor is used instead. `pickle` works like expected.

## 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: 1.4.0
 - OS (e.g., Linux): Ubuntu 18.04.3 LTS
 - How you installed PyTorch (`conda`, `pip`, source): conda
 - Build command you used (if compiling from source): N/A
 - Python version: 3.8.1
 - CUDA/cuDNN version: N/A
 - GPU models and configuration: N/A
 - Any other relevant information: Nope