Repository: pytorch
Issue: Cannot pickle torch model in version 2.1 when using `autograd.Function` as a model's attribute
Labels: needs reproduction,module: autograd,module: serialization,triaged,module: third_party

### 🐛 Describe the bug

I'm not able to serialise torch model in ray in torch 2.1.2 version. The model has an object of `autograd.Function` as an attribute. But I am able to serialise the same model in torch 1.13.1 version

I think it is a backwards incompatibility in torch 2.1.2

Here is a minimal reproducible snippet

```python
import ray
import torch
import torch.nn as nn

print(f"Checking serializability with torch version: {torch.__version__} and ray version: {ray.__version__}", end="\n\n\n")


class Operation(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return input

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output


class ModelWithOperation(nn.Module):
    def __init__(self, model: nn.Module):
        super().__init__()
        self.model = model
        self.operation = Operation.apply

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        output = self.model(input)
        return self.operation(output)


my_model = ModelWithOperation(model=nn.Linear(1, 1))
print(ray.util.inspect_serializability(my_model))
```


Output in torch version 2.1.2

```
Checking serializability with torch version: 2.1.2+cu121 and ray version: 2.4.0


================================================================================
Checking Serializability of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)
================================================================================
!!! FAIL serialization: cannot pickle 'torch._C.DispatchKeySet' object
    Serializing 'T_destination' ~T_destination...
    Serializing '_apply' <bound method Module._apply of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)>...
    !!! FAIL serialization: cannot pickle 'torch._C.DispatchKeySet' object
        Serializing '__func__' <function Module._apply at 0x7f3bd63b4a60>...
    WARNING: Did not find non-serializable object in <bound method Module._apply of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)>. This may be an oversight.
================================================================================
Variable: 

        FailTuple(_apply [obj=<bound method Module._apply of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)>, parent=ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)])

was found to be non-serializable. There may be multiple other undetected variables that were non-serializable. 
Consider either removing the instantiation/imports of these variables or moving the instantiation into the scope of the function/class. 
================================================================================
Check https://docs.ray.io/en/master/ray-core/objects/serialization.html#troubleshooting for more information.
If you have any suggestions on how to improve this error message, please reach out to the Ray developers on github.com/ray-project/ray/issues/
================================================================================
(False, {FailTuple(_apply [obj=<bound method Module._apply of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)>, parent=ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)])})
```

Output in torch version 1.13.1

```
Checking serializability with torch version: 1.13.1+cu117 and ray version: 2.4.0


================================================================================
Checking Serializability of ModelWithOperation(
  (model): Linear(in_features=1, out_features=1, bias=True)
)
================================================================================
(True, set())
```

### Versions

Collecting environment information...
PyTorch version: 1.13.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.28.1
Libc version: glibc-2.31

Python version: 3.9.5 (default, Nov 23 2021, 15:27:38)  [GCC 9.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1038-gcp-x86_64-with-glibc2.31
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
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

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   46 bits physical, 48 bits virtual
CPU(s):                          8
On-line CPU(s) list:             0-7
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           85
Model name:                      Intel(R) Xeon(R) CPU @ 2.00GHz
Stepping:                        3
CPU MHz:                         1999.999
BogoMIPS:                        3999.99
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       128 KiB
L1i cache:                       128 KiB
L2 cache:                        4 MiB
L3 cache:                        38.5 MiB

Versions of relevant libraries:
[pip3] flake8==3.9.2
[pip3] mypy==0.812
[pip3] mypy-extensions==0.4.4
[pip3] numpy==1.25.2
[pip3] pytorch-lightning==1.9.4
[pip3] torch==1.13.1
[pip3] torch-model-archiver==0.3.1
[pip3] torchmetrics==0.7.0
[pip3] torchserve==0.3.1
[pip3] triton==2.1.0
[conda] Could not collect

cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7 @mruberry @mikaylagawarecki