Cleanup in get_enum_from_fn (#8852)

Co-authored-by: Nicolas Hug <nh.nicolas.hug@gmail.com>

Author: GdoongMathew

torchvision/models/_api.py

@@ -7,7 +7,7 @@
 from functools import partial
 from inspect import signature
 from types import ModuleType
-from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Set, Type, TypeVar, Union
+from typing import Any, Callable, Dict, get_args, Iterable, List, Mapping, Optional, Set, Type, TypeVar, Union
 
 from torch import nn
 
@@ -168,14 +168,13 @@ def _get_enum_from_fn(fn: Callable) -> Type[WeightsEnum]:
  if "weights" not in sig.parameters:
  raise ValueError("The method is missing the 'weights' argument.")
 
- ann = signature(fn).parameters["weights"].annotation
+ ann = sig.parameters["weights"].annotation
  weights_enum = None
  if isinstance(ann, type) and issubclass(ann, WeightsEnum):
  weights_enum = ann
  else:
  # handle cases like Union[Optional, T]
- # TODO: Replace ann.__args__ with typing.get_args(ann) after python >= 3.8
- for t in ann.__args__: # type: ignore[union-attr]
+ for t in get_args(ann): # type: ignore[union-attr]
  if isinstance(t, type) and issubclass(t, WeightsEnum):
  weights_enum = t
  break