Repository: pytorch
Issue: [JIT] JIT does not support non-persistent buffers
Labels: oncall: jit

## 🐛 Bug

Models using JIT do not appear to respect the ``persistent=False`` option for buffers, and include all buffers in their ``state_dict`` regardless of the persistence setting. This behavior is not consistent with what the non-JIT model does and leads to the JIT and non-JIT models having incompatible ``state_dict``s (when using strict ``state_dict`` loading).

## To Reproduce

Steps to reproduce the behavior:

1. Run the following python script:
```python
import traceback

import torch
import torch.nn as nn


class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.register_buffer("buffer", torch.zeros(1), persistent=False)


def test_state_dict(model_base, jit: bool = False):
    model_jit = torch.jit.script(model_base)

    model = model_jit if jit else model_base

    saved_state_dict = model.state_dict()

    print("JIT:", jit)
    print("  state_dict:", saved_state_dict, "\n")

    model.load_state_dict(saved_state_dict, strict=True)

    try:
        model_base.load_state_dict(saved_state_dict, strict=True)
    except RuntimeError:
        traceback.print_exc()

    try:
        model_jit.load_state_dict(saved_state_dict, strict=True)
    except RuntimeError:
        traceback.print_exc()


print("PyTorch version:", torch.__version__, "\n")

base_model = Model()

test_state_dict(base_model, False)

print("\n----------------\n")

test_state_dict(base_model, True)

```
2. Observe (sample) output:
```
PyTorch version: 1.9.0

JIT: False
  state_dict: OrderedDict()

Traceback (most recent call last):
  File "mcve.py", line 32, in test_state_dict
    model_jit.load_state_dict(saved_state_dict, strict=True)
  File "/home/install/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1407, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for RecursiveScriptModule:
        Missing key(s) in state_dict: "buffer".

----------------

JIT: True
  state_dict: OrderedDict([('buffer', tensor([0.]))])

Traceback (most recent call last):
  File "mcve.py", line 27, in test_state_dict
    model_base.load_state_dict(saved_state_dict, strict=True)
  File "/home/install/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1407, in load_state_dict
    self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for Model:
        Unexpected key(s) in state_dict: "buffer".
```

## Expected behavior

The behavior of the JIT and non-JIT version of the model should be the same. It should be possible to perform seamless saving and loading of ``state_dict``s (even with strict loading) between JIT and non-JIT versions of the model.

## Environment

```
PyTorch version: 1.9.0
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Arch Linux (x86_64)
GCC version: (GCC) 11.1.0
Clang version: 12.0.1
CMake version: version 3.21.1
Libc version: glibc-2.17

Python version: 3.7.11 (default, Jul 27 2021, 14:32:16)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.13.8-arch1-1-x86_64-with-arch
Is CUDA available: True
CUDA runtime version: 11.4.48
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 2080 Ti
Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/libcudnn.so.8.2.2
/usr/lib/libcudnn_adv_infer.so.8.2.2
/usr/lib/libcudnn_adv_train.so.8.2.2
/usr/lib/libcudnn_cnn_infer.so.8.2.2
/usr/lib/libcudnn_cnn_train.so.8.2.2
/usr/lib/libcudnn_ops_infer.so.8.2.2
/usr/lib/libcudnn_ops_train.so.8.2.2
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.20.3
[pip3] torch==1.9.0
[pip3] torchvision==0.10.0
[conda] _tflow_select             2.3.0                       mkl
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.1.1               h6406543_8    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] mkl                       2021.3.0           h06a4308_520
[conda] mkl-service               2.4.0            py37h7f8727e_0
[conda] mkl_fft                   1.3.0            py37h42c9631_2
[conda] mkl_random                1.2.2            py37h51133e4_0
[conda] numpy                     1.18.0                   pypi_0    pypi
[conda] numpy-base                1.20.3           py37h74d4b33_0
[conda] pytorch                   1.9.0           py3.7_cuda11.1_cudnn8.0.5_0    pytorch
[conda] tensorflow                2.4.1           mkl_py37h2d14ff2_0
[conda] tensorflow-base           2.4.1           mkl_py37h43e0292_0
[conda] torchvision               0.10.0               py37_cu111    pytorch
```