Skip to content

Commit f670152

Browse files
check if value exists without iteration (#5225)
1 parent 03d1133 commit f670152

File tree

1 file changed

+2
-3
lines changed
  • torchvision/prototype/models

1 file changed

+2
-3
lines changed

torchvision/prototype/models/_api.py

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,9 +60,8 @@ def verify(cls, obj: Any) -> Any:
6060

6161
@classmethod
6262
def from_str(cls, value: str) -> "WeightsEnum":
63-
for k, v in cls.__members__.items():
64-
if k == value:
65-
return v
63+
if value in cls.__members__:
64+
return cls.__members__[value]
6665
raise ValueError(f"Invalid value {value} for enum {cls.__name__}.")
6766

6867
def get_state_dict(self, progress: bool) -> OrderedDict:

0 commit comments

Comments
 (0)