We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torchvision/models/_api.py`
Content:
```
1 import importlib
2 import inspect
3 import sys
4 from dataclasses import dataclass, fields
5 from inspect import signature
6 from types import ModuleType
7 from typing import Any, Callable, cast, Dict, List, Mapping, Optional, TypeVar, Union
8 
9 from torch import nn
10 
11 from torchvision._utils import StrEnum
12 
13 from .._internally_replaced_utils import load_state_dict_from_url
14 
15 
16 __all__ = ["WeightsEnum", "Weights", "get_model", "get_model_builder", "get_model_weights", "get_weight", "list_models"]
17 
18 
19 @dataclass
20 class Weights:
21     """
22     This class is used to group important attributes associated with the pre-trained weights.
23 
24     Args:
25         url (str): The location where we find the weights.
26         transforms (Callable): A callable that constructs the preprocessing method (or validation preset transforms)
27             needed to use the model. The reason we attach a constructor method rather than an already constructed
28             object is because the specific object might have memory and thus we want to delay initialization until
29             needed.
30         meta (Dict[str, Any]): Stores meta-data related to the weights of the model and its configuration. These can be
31             informative attributes (for example the number of parameters/flops, recipe link/methods used in training
32             etc), configuration parameters (for example the `num_classes`) needed to construct the model or important
33             meta-data (for example the `classes` of a classification model) needed to use the model.
34     """
35 
36     url: str
37     transforms: Callable
38     meta: Dict[str, Any]
39 
40 
41 class WeightsEnum(StrEnum):
42     """
43     This class is the parent class of all model weights. Each model building method receives an optional `weights`
44     parameter with its associated pre-trained weights. It inherits from `Enum` and its values should be of type
45     `Weights`.
46 
47     Args:
48         value (Weights): The data class entry with the weight information.
49     """
50 
51     def __init__(self, value: Weights):
52         self._value_ = value
53 
54     @classmethod
55     def verify(cls, obj: Any) -> Any:
56         if obj is not None:
57             if type(obj) is str:
58                 obj = cls.from_str(obj.replace(cls.__name__ + ".", ""))
59             elif not isinstance(obj, cls):
60                 raise TypeError(
61                     f"Invalid Weight class provided; expected {cls.__name__} but received {obj.__class__.__name__}."
62                 )
63         return obj
64 
65     def get_state_dict(self, progress: bool) -> Mapping[str, Any]:
66         return load_state_dict_from_url(self.url, progress=progress)
67 
68     def __repr__(self) -> str:
69         return f"{self.__class__.__name__}.{self._name_}"
70 
71     def __getattr__(self, name):
72         # Be able to fetch Weights attributes directly
73         for f in fields(Weights):
74             if f.name == name:
75                 return object.__getattribute__(self.value, name)
76         return super().__getattr__(name)
77 
78     def __deepcopy__(self, memodict=None):
79         return self
80 
81 
82 def get_weight(name: str) -> WeightsEnum:
83     """
84     Gets the weights enum value by its full name. Example: "ResNet50_Weights.IMAGENET1K_V1"
85 
86     .. betastatus:: function
87 
88     Args:
89         name (str): The name of the weight enum entry.
90 
91     Returns:
92         WeightsEnum: The requested weight enum.
93     """
94     try:
95         enum_name, value_name = name.split(".")
96     except ValueError:
97         raise ValueError(f"Invalid weight name provided: '{name}'.")
98 
99     base_module_name = ".".join(sys.modules[__name__].__name__.split(".")[:-1])
100     base_module = importlib.import_module(base_module_name)
101     model_modules = [base_module] + [
102         x[1] for x in inspect.getmembers(base_module, inspect.ismodule) if x[1].__file__.endswith("__init__.py")
103     ]
104 
105     weights_enum = None
106     for m in model_modules:
107         potential_class = m.__dict__.get(enum_name, None)
108         if potential_class is not None and issubclass(potential_class, WeightsEnum):
109             weights_enum = potential_class
110             break
111 
112     if weights_enum is None:
113         raise ValueError(f"The weight enum '{enum_name}' for the specific method couldn't be retrieved.")
114 
115     return weights_enum.from_str(value_name)
116 
117 
118 def get_model_weights(name: Union[Callable, str]) -> WeightsEnum:
119     """
120     Returns the weights enum class associated to the given model.
121 
122     .. betastatus:: function
123 
124     Args:
125         name (callable or str): The model builder function or the name under which it is registered.
126 
127     Returns:
128         weights_enum (WeightsEnum): The weights enum class associated with the model.
129     """
130     model = get_model_builder(name) if isinstance(name, str) else name
131     return _get_enum_from_fn(model)
132 
133 
134 def _get_enum_from_fn(fn: Callable) -> WeightsEnum:
135     """
136     Internal method that gets the weight enum of a specific model builder method.
137 
138     Args:
139         fn (Callable): The builder method used to create the model.
140     Returns:
141         WeightsEnum: The requested weight enum.
142     """
143     sig = signature(fn)
144     if "weights" not in sig.parameters:
145         raise ValueError("The method is missing the 'weights' argument.")
146 
147     ann = signature(fn).parameters["weights"].annotation
148     weights_enum = None
149     if isinstance(ann, type) and issubclass(ann, WeightsEnum):
150         weights_enum = ann
151     else:
152         # handle cases like Union[Optional, T]
153         # TODO: Replace ann.__args__ with typing.get_args(ann) after python >= 3.8
154         for t in ann.__args__:  # type: ignore[union-attr]
155             if isinstance(t, type) and issubclass(t, WeightsEnum):
156                 weights_enum = t
157                 break
158 
159     if weights_enum is None:
160         raise ValueError(
161             "The WeightsEnum class for the specific method couldn't be retrieved. Make sure the typing info is correct."
162         )
163 
164     return cast(WeightsEnum, weights_enum)
165 
166 
167 M = TypeVar("M", bound=nn.Module)
168 
169 BUILTIN_MODELS = {}
170 
171 
172 def register_model(name: Optional[str] = None) -> Callable[[Callable[..., M]], Callable[..., M]]:
173     def wrapper(fn: Callable[..., M]) -> Callable[..., M]:
174         key = name if name is not None else fn.__name__
175         if key in BUILTIN_MODELS:
176             raise ValueError(f"An entry is already registered under the name '{key}'.")
177         BUILTIN_MODELS[key] = fn
178         return fn
179 
180     return wrapper
181 
182 
183 def list_models(module: Optional[ModuleType] = None) -> List[str]:
184     """
185     Returns a list with the names of registered models.
186 
187     .. betastatus:: function
188 
189     Args:
190         module (ModuleType, optional): The module from which we want to extract the available models.
191 
192     Returns:
193         models (list): A list with the names of available models.
194     """
195     models = [
196         k for k, v in BUILTIN_MODELS.items() if module is None or v.__module__.rsplit(".", 1)[0] == module.__name__
197     ]
198     return sorted(models)
199 
200 
201 def get_model_builder(name: str) -> Callable[..., nn.Module]:
202     """
203     Gets the model name and returns the model builder method.
204 
205     .. betastatus:: function
206 
207     Args:
208         name (str): The name under which the model is registered.
209 
210     Returns:
211         fn (Callable): The model builder method.
212     """
213     name = name.lower()
214     try:
215         fn = BUILTIN_MODELS[name]
216     except KeyError:
217         raise ValueError(f"Unknown model {name}")
218     return fn
219 
220 
221 def get_model(name: str, **config: Any) -> nn.Module:
222     """
223     Gets the model name and configuration and returns an instantiated model.
224 
225     .. betastatus:: function
226 
227     Args:
228         name (str): The name under which the model is registered.
229         **config (Any): parameters passed to the model builder method.
230 
231     Returns:
232         model (nn.Module): The initialized model.
233     """
234     fn = get_model_builder(name)
235     return fn(**config)
236
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
                 return object.__getattribute__(self.value, name)
         return super().__getattr__(name)
 
-    def __deepcopy__(self, memodict=None):
-        return self
-
 
 def get_weight(name: str) -> WeightsEnum:
     """

{"golden_diff": "diff --git a/torchvision/models/_api.py b/torchvision/models/_api.py\n--- a/torchvision/models/_api.py\n+++ b/torchvision/models/_api.py\n@@ -2,6 +2,7 @@\n import inspect\n import sys\n from dataclasses import dataclass, fields\n+from functools import partial\n from inspect import signature\n from types import ModuleType\n from typing import Any, Callable, cast, Dict, List, Mapping, Optional, TypeVar, Union\n@@ -37,6 +38,32 @@\n     transforms: Callable\n     meta: Dict[str, Any]\n \n+    def __eq__(self, other: Any) -> bool:\n+        # We need this custom implementation for correct deep-copy and deserialization behavior.\n+        # TL;DR: After the definition of an enum, creating a new instance, i.e. by deep-copying or deserializing it,\n+        # involves an equality check against the defined members. Unfortunately, the `transforms` attribute is often\n+        # defined with `functools.partial` and `fn = partial(...); assert deepcopy(fn) != fn`. Without custom handling\n+        # for it, the check against the defined members would fail and effectively prevent the weights from being\n+        # deep-copied or deserialized.\n+        # See https://github.com/pytorch/vision/pull/7107 for details.\n+        if not isinstance(other, Weights):\n+            return NotImplemented\n+\n+        if self.url != other.url:\n+            return False\n+\n+        if self.meta != other.meta:\n+            return False\n+\n+        if isinstance(self.transforms, partial) and isinstance(other.transforms, partial):\n+            return (\n+                self.transforms.func == other.transforms.func\n+                and self.transforms.args == other.transforms.args\n+                and self.transforms.keywords == other.transforms.keywords\n+            )\n+        else:\n+            return self.transforms == other.transforms\n+\n \n class WeightsEnum(StrEnum):\n     \"\"\"\n@@ -75,9 +102,6 @@\n                 return object.__getattribute__(self.value, name)\n         return super().__getattr__(name)\n \n-    def __deepcopy__(self, memodict=None):\n-        return self\n-\n \n def get_weight(name: str) -> WeightsEnum:\n     \"\"\"\n", "issue": "Weights enums cannot be pickled\n### \ud83d\udc1b Describe the bug\r\n\r\nTrying 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.\r\n\r\n```python\r\nimport pickle\r\n\r\nfrom torchvision.models import ResNet50_Weights\r\n\r\n\r\ndef main():\r\n    w = ResNet50_Weights.DEFAULT\r\n    kek = pickle.dumps(w)\r\n    lol = pickle.loads(kek)\r\n    print(lol)\r\n\r\n\r\nif __name__ == '__main__':\r\n    main()\r\n```\r\n\r\nThis 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\r\n\r\n```python\r\nimport functools\r\nfrom torchvision.models._api import Weights\r\n\r\ndef new_eq(self, other):\r\n        if not isinstance(other, type(self)):\r\n            return False\r\n        if self.meta != other.meta or self.url != other.url:\r\n            return False\r\n        if not isinstance(self.transforms, functools.partial) or not isinstance(other.transforms, functools.partial):\r\n            return self.transforms == other.transforms\r\n        return all(getattr(self.transforms, a) == getattr(other.transforms, a) for a in [\"func\", \"args\", \"keywords\", \"__dict__\"])\r\n\r\nWeights.__eq__ = new_eq\r\n```\r\n\r\n### Versions\r\n\r\n```\r\nCollecting environment information...\r\nPyTorch version: 1.12.1\r\nIs debug build: False\r\nCUDA used to build PyTorch: None\r\nROCM used to build PyTorch: N/A\r\n\r\nOS: macOS 13.0.1 (arm64)\r\nGCC version: Could not collect\r\nClang version: 14.0.0 (clang-1400.0.29.202)\r\nCMake version: version 3.25.0\r\nLibc version: N/A\r\n\r\nPython version: 3.9.13 | packaged by conda-forge | (main, May 27 2022, 17:00:33)  [Clang 13.0.1 ] (64-bit runtime)\r\nPython platform: macOS-13.0.1-arm64-arm-64bit\r\nIs CUDA available: False\r\nCUDA runtime version: No CUDA\r\nCUDA_MODULE_LOADING set to: N/A\r\nGPU models and configuration: No CUDA\r\nNvidia driver version: No CUDA\r\ncuDNN version: No CUDA\r\nHIP runtime version: N/A\r\nMIOpen runtime version: N/A\r\nIs XNNPACK available: True\r\n\r\nVersions of relevant libraries:\r\n[pip3] mypy==0.971\r\n[pip3] mypy-extensions==0.4.3\r\n[pip3] numpy==1.22.4\r\n[pip3] torch==1.12.1\r\n[pip3] torchvision==0.13.1\r\n[conda] numpy                     1.22.4           py39h7df2422_0    conda-forge\r\n[conda] torch                     1.12.1                   pypi_0    pypi\r\n[conda] torchvision               0.13.1                   pypi_0    pypi\r\n```\n", "before_files": [{"content": "import importlib\nimport inspect\nimport sys\nfrom dataclasses import dataclass, fields\nfrom inspect import signature\nfrom types import ModuleType\nfrom typing import Any, Callable, cast, Dict, List, Mapping, Optional, TypeVar, Union\n\nfrom torch import nn\n\nfrom torchvision._utils import StrEnum\n\nfrom .._internally_replaced_utils import load_state_dict_from_url\n\n\n__all__ = [\"WeightsEnum\", \"Weights\", \"get_model\", \"get_model_builder\", \"get_model_weights\", \"get_weight\", \"list_models\"]\n\n\n@dataclass\nclass Weights:\n    \"\"\"\n    This class is used to group important attributes associated with the pre-trained weights.\n\n    Args:\n        url (str): The location where we find the weights.\n        transforms (Callable): A callable that constructs the preprocessing method (or validation preset transforms)\n            needed to use the model. The reason we attach a constructor method rather than an already constructed\n            object is because the specific object might have memory and thus we want to delay initialization until\n            needed.\n        meta (Dict[str, Any]): Stores meta-data related to the weights of the model and its configuration. These can be\n            informative attributes (for example the number of parameters/flops, recipe link/methods used in training\n            etc), configuration parameters (for example the `num_classes`) needed to construct the model or important\n            meta-data (for example the `classes` of a classification model) needed to use the model.\n    \"\"\"\n\n    url: str\n    transforms: Callable\n    meta: Dict[str, Any]\n\n\nclass WeightsEnum(StrEnum):\n    \"\"\"\n    This class is the parent class of all model weights. Each model building method receives an optional `weights`\n    parameter with its associated pre-trained weights. It inherits from `Enum` and its values should be of type\n    `Weights`.\n\n    Args:\n        value (Weights): The data class entry with the weight information.\n    \"\"\"\n\n    def __init__(self, value: Weights):\n        self._value_ = value\n\n    @classmethod\n    def verify(cls, obj: Any) -> Any:\n        if obj is not None:\n            if type(obj) is str:\n                obj = cls.from_str(obj.replace(cls.__name__ + \".\", \"\"))\n            elif not isinstance(obj, cls):\n                raise TypeError(\n                    f\"Invalid Weight class provided; expected {cls.__name__} but received {obj.__class__.__name__}.\"\n                )\n        return obj\n\n    def get_state_dict(self, progress: bool) -> Mapping[str, Any]:\n        return load_state_dict_from_url(self.url, progress=progress)\n\n    def __repr__(self) -> str:\n        return f\"{self.__class__.__name__}.{self._name_}\"\n\n    def __getattr__(self, name):\n        # Be able to fetch Weights attributes directly\n        for f in fields(Weights):\n            if f.name == name:\n                return object.__getattribute__(self.value, name)\n        return super().__getattr__(name)\n\n    def __deepcopy__(self, memodict=None):\n        return self\n\n\ndef get_weight(name: str) -> WeightsEnum:\n    \"\"\"\n    Gets the weights enum value by its full name. Example: \"ResNet50_Weights.IMAGENET1K_V1\"\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name of the weight enum entry.\n\n    Returns:\n        WeightsEnum: The requested weight enum.\n    \"\"\"\n    try:\n        enum_name, value_name = name.split(\".\")\n    except ValueError:\n        raise ValueError(f\"Invalid weight name provided: '{name}'.\")\n\n    base_module_name = \".\".join(sys.modules[__name__].__name__.split(\".\")[:-1])\n    base_module = importlib.import_module(base_module_name)\n    model_modules = [base_module] + [\n        x[1] for x in inspect.getmembers(base_module, inspect.ismodule) if x[1].__file__.endswith(\"__init__.py\")\n    ]\n\n    weights_enum = None\n    for m in model_modules:\n        potential_class = m.__dict__.get(enum_name, None)\n        if potential_class is not None and issubclass(potential_class, WeightsEnum):\n            weights_enum = potential_class\n            break\n\n    if weights_enum is None:\n        raise ValueError(f\"The weight enum '{enum_name}' for the specific method couldn't be retrieved.\")\n\n    return weights_enum.from_str(value_name)\n\n\ndef get_model_weights(name: Union[Callable, str]) -> WeightsEnum:\n    \"\"\"\n    Returns the weights enum class associated to the given model.\n\n    .. betastatus:: function\n\n    Args:\n        name (callable or str): The model builder function or the name under which it is registered.\n\n    Returns:\n        weights_enum (WeightsEnum): The weights enum class associated with the model.\n    \"\"\"\n    model = get_model_builder(name) if isinstance(name, str) else name\n    return _get_enum_from_fn(model)\n\n\ndef _get_enum_from_fn(fn: Callable) -> WeightsEnum:\n    \"\"\"\n    Internal method that gets the weight enum of a specific model builder method.\n\n    Args:\n        fn (Callable): The builder method used to create the model.\n    Returns:\n        WeightsEnum: The requested weight enum.\n    \"\"\"\n    sig = signature(fn)\n    if \"weights\" not in sig.parameters:\n        raise ValueError(\"The method is missing the 'weights' argument.\")\n\n    ann = signature(fn).parameters[\"weights\"].annotation\n    weights_enum = None\n    if isinstance(ann, type) and issubclass(ann, WeightsEnum):\n        weights_enum = ann\n    else:\n        # handle cases like Union[Optional, T]\n        # TODO: Replace ann.__args__ with typing.get_args(ann) after python >= 3.8\n        for t in ann.__args__:  # type: ignore[union-attr]\n            if isinstance(t, type) and issubclass(t, WeightsEnum):\n                weights_enum = t\n                break\n\n    if weights_enum is None:\n        raise ValueError(\n            \"The WeightsEnum class for the specific method couldn't be retrieved. Make sure the typing info is correct.\"\n        )\n\n    return cast(WeightsEnum, weights_enum)\n\n\nM = TypeVar(\"M\", bound=nn.Module)\n\nBUILTIN_MODELS = {}\n\n\ndef register_model(name: Optional[str] = None) -> Callable[[Callable[..., M]], Callable[..., M]]:\n    def wrapper(fn: Callable[..., M]) -> Callable[..., M]:\n        key = name if name is not None else fn.__name__\n        if key in BUILTIN_MODELS:\n            raise ValueError(f\"An entry is already registered under the name '{key}'.\")\n        BUILTIN_MODELS[key] = fn\n        return fn\n\n    return wrapper\n\n\ndef list_models(module: Optional[ModuleType] = None) -> List[str]:\n    \"\"\"\n    Returns a list with the names of registered models.\n\n    .. betastatus:: function\n\n    Args:\n        module (ModuleType, optional): The module from which we want to extract the available models.\n\n    Returns:\n        models (list): A list with the names of available models.\n    \"\"\"\n    models = [\n        k for k, v in BUILTIN_MODELS.items() if module is None or v.__module__.rsplit(\".\", 1)[0] == module.__name__\n    ]\n    return sorted(models)\n\n\ndef get_model_builder(name: str) -> Callable[..., nn.Module]:\n    \"\"\"\n    Gets the model name and returns the model builder method.\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name under which the model is registered.\n\n    Returns:\n        fn (Callable): The model builder method.\n    \"\"\"\n    name = name.lower()\n    try:\n        fn = BUILTIN_MODELS[name]\n    except KeyError:\n        raise ValueError(f\"Unknown model {name}\")\n    return fn\n\n\ndef get_model(name: str, **config: Any) -> nn.Module:\n    \"\"\"\n    Gets the model name and configuration and returns an instantiated model.\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name under which the model is registered.\n        **config (Any): parameters passed to the model builder method.\n\n    Returns:\n        model (nn.Module): The initialized model.\n    \"\"\"\n    fn = get_model_builder(name)\n    return fn(**config)\n", "path": "torchvision/models/_api.py"}], "after_files": [{"content": "import importlib\nimport inspect\nimport sys\nfrom dataclasses import dataclass, fields\nfrom functools import partial\nfrom inspect import signature\nfrom types import ModuleType\nfrom typing import Any, Callable, cast, Dict, List, Mapping, Optional, TypeVar, Union\n\nfrom torch import nn\n\nfrom torchvision._utils import StrEnum\n\nfrom .._internally_replaced_utils import load_state_dict_from_url\n\n\n__all__ = [\"WeightsEnum\", \"Weights\", \"get_model\", \"get_model_builder\", \"get_model_weights\", \"get_weight\", \"list_models\"]\n\n\n@dataclass\nclass Weights:\n    \"\"\"\n    This class is used to group important attributes associated with the pre-trained weights.\n\n    Args:\n        url (str): The location where we find the weights.\n        transforms (Callable): A callable that constructs the preprocessing method (or validation preset transforms)\n            needed to use the model. The reason we attach a constructor method rather than an already constructed\n            object is because the specific object might have memory and thus we want to delay initialization until\n            needed.\n        meta (Dict[str, Any]): Stores meta-data related to the weights of the model and its configuration. These can be\n            informative attributes (for example the number of parameters/flops, recipe link/methods used in training\n            etc), configuration parameters (for example the `num_classes`) needed to construct the model or important\n            meta-data (for example the `classes` of a classification model) needed to use the model.\n    \"\"\"\n\n    url: str\n    transforms: Callable\n    meta: Dict[str, Any]\n\n    def __eq__(self, other: Any) -> bool:\n        # We need this custom implementation for correct deep-copy and deserialization behavior.\n        # TL;DR: After the definition of an enum, creating a new instance, i.e. by deep-copying or deserializing it,\n        # involves an equality check against the defined members. Unfortunately, the `transforms` attribute is often\n        # defined with `functools.partial` and `fn = partial(...); assert deepcopy(fn) != fn`. Without custom handling\n        # for it, the check against the defined members would fail and effectively prevent the weights from being\n        # deep-copied or deserialized.\n        # See https://github.com/pytorch/vision/pull/7107 for details.\n        if not isinstance(other, Weights):\n            return NotImplemented\n\n        if self.url != other.url:\n            return False\n\n        if self.meta != other.meta:\n            return False\n\n        if isinstance(self.transforms, partial) and isinstance(other.transforms, partial):\n            return (\n                self.transforms.func == other.transforms.func\n                and self.transforms.args == other.transforms.args\n                and self.transforms.keywords == other.transforms.keywords\n            )\n        else:\n            return self.transforms == other.transforms\n\n\nclass WeightsEnum(StrEnum):\n    \"\"\"\n    This class is the parent class of all model weights. Each model building method receives an optional `weights`\n    parameter with its associated pre-trained weights. It inherits from `Enum` and its values should be of type\n    `Weights`.\n\n    Args:\n        value (Weights): The data class entry with the weight information.\n    \"\"\"\n\n    def __init__(self, value: Weights):\n        self._value_ = value\n\n    @classmethod\n    def verify(cls, obj: Any) -> Any:\n        if obj is not None:\n            if type(obj) is str:\n                obj = cls.from_str(obj.replace(cls.__name__ + \".\", \"\"))\n            elif not isinstance(obj, cls):\n                raise TypeError(\n                    f\"Invalid Weight class provided; expected {cls.__name__} but received {obj.__class__.__name__}.\"\n                )\n        return obj\n\n    def get_state_dict(self, progress: bool) -> Mapping[str, Any]:\n        return load_state_dict_from_url(self.url, progress=progress)\n\n    def __repr__(self) -> str:\n        return f\"{self.__class__.__name__}.{self._name_}\"\n\n    def __getattr__(self, name):\n        # Be able to fetch Weights attributes directly\n        for f in fields(Weights):\n            if f.name == name:\n                return object.__getattribute__(self.value, name)\n        return super().__getattr__(name)\n\n\ndef get_weight(name: str) -> WeightsEnum:\n    \"\"\"\n    Gets the weights enum value by its full name. Example: \"ResNet50_Weights.IMAGENET1K_V1\"\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name of the weight enum entry.\n\n    Returns:\n        WeightsEnum: The requested weight enum.\n    \"\"\"\n    try:\n        enum_name, value_name = name.split(\".\")\n    except ValueError:\n        raise ValueError(f\"Invalid weight name provided: '{name}'.\")\n\n    base_module_name = \".\".join(sys.modules[__name__].__name__.split(\".\")[:-1])\n    base_module = importlib.import_module(base_module_name)\n    model_modules = [base_module] + [\n        x[1] for x in inspect.getmembers(base_module, inspect.ismodule) if x[1].__file__.endswith(\"__init__.py\")\n    ]\n\n    weights_enum = None\n    for m in model_modules:\n        potential_class = m.__dict__.get(enum_name, None)\n        if potential_class is not None and issubclass(potential_class, WeightsEnum):\n            weights_enum = potential_class\n            break\n\n    if weights_enum is None:\n        raise ValueError(f\"The weight enum '{enum_name}' for the specific method couldn't be retrieved.\")\n\n    return weights_enum.from_str(value_name)\n\n\ndef get_model_weights(name: Union[Callable, str]) -> WeightsEnum:\n    \"\"\"\n    Returns the weights enum class associated to the given model.\n\n    .. betastatus:: function\n\n    Args:\n        name (callable or str): The model builder function or the name under which it is registered.\n\n    Returns:\n        weights_enum (WeightsEnum): The weights enum class associated with the model.\n    \"\"\"\n    model = get_model_builder(name) if isinstance(name, str) else name\n    return _get_enum_from_fn(model)\n\n\ndef _get_enum_from_fn(fn: Callable) -> WeightsEnum:\n    \"\"\"\n    Internal method that gets the weight enum of a specific model builder method.\n\n    Args:\n        fn (Callable): The builder method used to create the model.\n    Returns:\n        WeightsEnum: The requested weight enum.\n    \"\"\"\n    sig = signature(fn)\n    if \"weights\" not in sig.parameters:\n        raise ValueError(\"The method is missing the 'weights' argument.\")\n\n    ann = signature(fn).parameters[\"weights\"].annotation\n    weights_enum = None\n    if isinstance(ann, type) and issubclass(ann, WeightsEnum):\n        weights_enum = ann\n    else:\n        # handle cases like Union[Optional, T]\n        # TODO: Replace ann.__args__ with typing.get_args(ann) after python >= 3.8\n        for t in ann.__args__:  # type: ignore[union-attr]\n            if isinstance(t, type) and issubclass(t, WeightsEnum):\n                weights_enum = t\n                break\n\n    if weights_enum is None:\n        raise ValueError(\n            \"The WeightsEnum class for the specific method couldn't be retrieved. Make sure the typing info is correct.\"\n        )\n\n    return cast(WeightsEnum, weights_enum)\n\n\nM = TypeVar(\"M\", bound=nn.Module)\n\nBUILTIN_MODELS = {}\n\n\ndef register_model(name: Optional[str] = None) -> Callable[[Callable[..., M]], Callable[..., M]]:\n    def wrapper(fn: Callable[..., M]) -> Callable[..., M]:\n        key = name if name is not None else fn.__name__\n        if key in BUILTIN_MODELS:\n            raise ValueError(f\"An entry is already registered under the name '{key}'.\")\n        BUILTIN_MODELS[key] = fn\n        return fn\n\n    return wrapper\n\n\ndef list_models(module: Optional[ModuleType] = None) -> List[str]:\n    \"\"\"\n    Returns a list with the names of registered models.\n\n    .. betastatus:: function\n\n    Args:\n        module (ModuleType, optional): The module from which we want to extract the available models.\n\n    Returns:\n        models (list): A list with the names of available models.\n    \"\"\"\n    models = [\n        k for k, v in BUILTIN_MODELS.items() if module is None or v.__module__.rsplit(\".\", 1)[0] == module.__name__\n    ]\n    return sorted(models)\n\n\ndef get_model_builder(name: str) -> Callable[..., nn.Module]:\n    \"\"\"\n    Gets the model name and returns the model builder method.\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name under which the model is registered.\n\n    Returns:\n        fn (Callable): The model builder method.\n    \"\"\"\n    name = name.lower()\n    try:\n        fn = BUILTIN_MODELS[name]\n    except KeyError:\n        raise ValueError(f\"Unknown model {name}\")\n    return fn\n\n\ndef get_model(name: str, **config: Any) -> nn.Module:\n    \"\"\"\n    Gets the model name and configuration and returns an instantiated model.\n\n    .. betastatus:: function\n\n    Args:\n        name (str): The name under which the model is registered.\n        **config (Any): parameters passed to the model builder method.\n\n    Returns:\n        model (nn.Module): The initialized model.\n    \"\"\"\n    fn = get_model_builder(name)\n    return fn(**config)\n", "path": "torchvision/models/_api.py"}]}