Repository: pytorch/vision

Issue:
Weights enums cannot be pickled
### 🐛 Describe the bug

Trying to unplickle pickled weights enum object yields an error inside enum implementation (it cannot match constructed object with existing enum values). Here is reproducible example.

```python
import pickle

from torchvision.models import ResNet50_Weights


def main():
    w = ResNet50_Weights.DEFAULT
    kek = pickle.dumps(w)
    lol = pickle.loads(kek)
    print(lol)


if __name__ == '__main__':
    main()
```

This happens because internally enum tries to match newly constructed enum value with declared values. And this fails because `Weights` dataclass has `transforms` field which is a callable. And if that callable is a `functools.partial`, it is technically a different object, even though it was constructed from the same original function with the same partial arguments. This snippet fixes the error, so you should get the idea for probable fix

```python
import functools
from torchvision.models._api import Weights

def new_eq(self, other):
        if not isinstance(other, type(self)):
            return False
        if self.meta != other.meta or self.url != other.url:
            return False
        if not isinstance(self.transforms, functools.partial) or not isinstance(other.transforms, functools.partial):
            return self.transforms == other.transforms
        return all(getattr(self.transforms, a) == getattr(other.transforms, a) for a in ["func", "args", "keywords", "__dict__"])

Weights.__eq__ = new_eq
```

### Versions

```
Collecting environment information...
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.0.1 (arm64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.202)
CMake version: version 3.25.0
Libc version: N/A

Python version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 17:00:33)  [Clang 13.0.1 ] (64-bit runtime)
Python platform: macOS-13.0.1-arm64-arm-64bit
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

Versions of relevant libraries:
[pip3] mypy==0.971
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.22.4
[pip3] torch==1.12.1
[pip3] torchvision==0.13.1
[conda] numpy                     1.22.4           py39h7df2422_0    conda-forge
[conda] torch                     1.12.1                   pypi_0    pypi
[conda] torchvision               0.13.1                   pypi_0    pypi
```


Discussion:

Resolving patch:
diff --git a/torchvision/models/_api.py b/torchvision/models/_api.py
--- a/torchvision/models/_api.py
+++ b/torchvision/models/_api.py
@@ -2,6 +2,7 @@
 import inspect
 import sys
 from dataclasses import dataclass, fields
+from functools import partial
 from inspect import signature
 from types import ModuleType
 from typing import Any, Callable, cast, Dict, List, Mapping, Optional, TypeVar, Union
@@ -37,6 +38,32 @@ class Weights:
     transforms: Callable
     meta: Dict[str, Any]
 
+    def __eq__(self, other: Any) -> bool:
+        # We need this custom implementation for correct deep-copy and deserialization behavior.
+        # TL;DR: After the definition of an enum, creating a new instance, i.e. by deep-copying or deserializing it,
+        # involves an equality check against the defined members. Unfortunately, the `transforms` attribute is often
+        # defined with `functools.partial` and `fn = partial(...); assert deepcopy(fn) != fn`. Without custom handling
+        # for it, the check against the defined members would fail and effectively prevent the weights from being
+        # deep-copied or deserialized.
+        # See https://github.com/pytorch/vision/pull/7107 for details.
+        if not isinstance(other, Weights):
+            return NotImplemented
+
+        if self.url != other.url:
+            return False
+
+        if self.meta != other.meta:
+            return False
+
+        if isinstance(self.transforms, partial) and isinstance(other.transforms, partial):
+            return (
+                self.transforms.func == other.transforms.func
+                and self.transforms.args == other.transforms.args
+                and self.transforms.keywords == other.transforms.keywords
+            )
+        else:
+            return self.transforms == other.transforms
+
 
 class WeightsEnum(StrEnum):
     """
@@ -75,9 +102,6 @@ def __getattr__(self, name):
                 return object.__getattribute__(self.value, name)
         return super().__getattr__(name)
 
-    def __deepcopy__(self, memodict=None):
-        return self
-
 
 def get_weight(name: str) -> WeightsEnum:
     """


Test patch:
diff --git a/test/test_extended_models.py b/test/test_extended_models.py
--- a/test/test_extended_models.py
+++ b/test/test_extended_models.py
@@ -1,5 +1,6 @@
 import copy
 import os
+import pickle
 
 import pytest
 import test_models as TM
@@ -73,10 +74,32 @@ def test_get_model_weights(name, weight):
     ],
 )
 def test_weights_copyable(copy_fn, name):
-    model_weights = models.get_model_weights(name)
-    for weights in list(model_weights):
-        copied_weights = copy_fn(weights)
-        assert copied_weights is weights
+    for weights in list(models.get_model_weights(name)):
+        # It is somewhat surprising that (deep-)copying is an identity operation here, but this is the default behavior
+        # of enums: https://docs.python.org/3/howto/enum.html#enum-members-aka-instances
+        # Checking for equality, i.e. `==`, is sufficient (and even preferable) for our use case, should we need to drop
+        # support for the identity operation in the future.
+        assert copy_fn(weights) is weights
+
+
+@pytest.mark.parametrize(
+    "name",
+    [
+        "resnet50",
+        "retinanet_resnet50_fpn_v2",
+        "raft_large",
+        "quantized_resnet50",
+        "lraspp_mobilenet_v3_large",
+        "mvit_v1_b",
+    ],
+)
+def test_weights_deserializable(name):
+    for weights in list(models.get_model_weights(name)):
+        # It is somewhat surprising that deserialization is an identity operation here, but this is the default behavior
+        # of enums: https://docs.python.org/3/howto/enum.html#enum-members-aka-instances
+        # Checking for equality, i.e. `==`, is sufficient (and even preferable) for our use case, should we need to drop
+        # support for the identity operation in the future.
+        assert pickle.loads(pickle.dumps(weights)) is weights
 
 
 @pytest.mark.parametrize(